Skip to content

fix(pi): retire interrupted model selections #32450

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
StiensWout wants to merge 3 commits into from
+179 −66
Closed

fix(pi): retire interrupted model selections #32450

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
01c58cd
fix(pi): retire uncertain model selections
StiensWout Nov 7, 2026
a7ee627
test(pi): trim selection retirement tests
StiensWout Nov 7, 2026
c91e04b
Merge remote-tracking branch 'github/main' into pr-32450
maria-rcks Nov 10, 2026
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
109 changes: 106 additions & 3 deletions packages/provider-pi/src/server/adapter.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ import * as Queue from "effect/Queue";
import * as Schema from "effect/Schema";
import * as Sink from "effect/Sink";
import * as Stream from "effect/Stream";
import * as TestClock from "effect/testing/TestClock";
import * as ChildProcess from "effect/process/ChildProcess";
import * as ChildProcessSpawner from "effect/process/ChildProcessSpawner";

Expand Down Expand Up @@ -168,7 +169,7 @@ interface FakePi {
readonly failNextState: () => void;
readonly failNextThinking: () => void;
readonly deferNextLifecycle: (type: "switch_session" | "new_session" | "fork" | "abort") => void;
readonly deferNextModelSelection: () => void;
readonly deferNextModelSelection: (type?: "set_model" | "set_thinking_level") => void;
readonly queueModels: (models: ReadonlyArray<unknown>) => void;
readonly vetoNextNewSession: () => void;
readonly allRequests: () => ReadonlyArray<PiRpcRecord>;
Expand Down Expand Up @@ -452,8 +453,8 @@ const makeFakePi: Effect.Effect<FakePi> = Effect.gen(function* () {
deferNextLifecycle: (type) => {
deferredLifecycle = type;
},
deferNextModelSelection: () => {
deferredLifecycle = "set_model";
deferNextModelSelection: (type = "set_model") => {
deferredLifecycle = type;
},
queueModels: (value) => {
models = value;
Expand Down Expand Up @@ -2159,6 +2160,108 @@ describe("PiAdapterV2", () => {
}).pipe(Effect.scoped, Effect.provide(layerTest)),
);

it.effect.each([
{ operation: "set_model", restore: false, failure: "timeout" },
{ operation: "set_model", restore: false, failure: "interruption" },
{ operation: "set_thinking_level", restore: false, failure: "timeout" },
{ operation: "set_thinking_level", restore: false, failure: "interruption" },
{ operation: "set_model", restore: true, failure: "timeout" },
{ operation: "set_model", restore: true, failure: "interruption" },
{ operation: "set_thinking_level", restore: true, failure: "timeout" },
{ operation: "set_thinking_level", restore: true, failure: "interruption" },
] as const)(
"retires selection after $failure of $operation (baseline restore: $restore)",
({ operation, restore, failure }) =>
Effect.gen(function* () {
const fake = yield* makeFakePi;
fake.queueState({ model: { provider: "anthropic", id: "baseline" } });
const { runtime, takeEvent } = yield* openRuntime(fake);
const providerThread = yield* runtime.ensureThread({
threadId: THREAD_ID,
modelSelection: modelSelection("default"),
runtimePolicy,
});
if (restore) {
yield* startTurn(runtime, providerThread, "default", [], "First", {
...modelSelection("anthropic/model-a"),
options: [{ id: "thinking", value: "low" }],
});
yield* fake.takeRequest("set_model");
yield* fake.takeRequest("set_thinking_level");
yield* fake.takeRequest("prompt");
yield* fake.emit({ type: "agent_start" });
yield* fake.emit({ type: "agent_end", messages: [], willRetry: false });
yield* fake.emit({ type: "agent_settled" });
yield* takeEvent((event) => event.type === "turn.terminal");
}
fake.deferNextModelSelection(operation);
const selection: ModelSelection = restore
? modelSelection("default")
: {
...modelSelection("anthropic/model-a"),
options: [{ id: "thinking", value: "low" }],
};
const starting = yield* startTurn(
runtime,
providerThread,
"default",
[],
"Selection",
selection,
2,
).pipe(Effect.exit, Effect.forkChild);
const pending = yield* fake.takeRequest(operation);
if (failure === "timeout") {
yield* TestClock.adjust("15 seconds");
const result = yield* Fiber.join(starting);
assert.equal(result._tag, "Failure");
if (Exit.isFailure(result)) assert.match(Cause.pretty(result.cause), /PiRpcTimeoutError/);
} else {
yield* Fiber.interrupt(starting);
}
// A late acknowledgement cannot make the retired connection usable again.
yield* fake.emit({ type: "response", id: pending.id, command: operation, success: true });
const retry = yield* startTurn(runtime, providerThread, "anthropic/model-a").pipe(
Effect.exit,
);
assert.equal(retry._tag, "Failure");
assert.deepEqual(fake.promptModels, restore ? ["anthropic/model-a"] : []);
yield* runtime.resumeThread({ providerThread }).pipe(Effect.flip);
}).pipe(Effect.scoped, Effect.provide(layerTest)),
);

it.effect("keeps the runtime usable after Pi rejects a model selection", () =>
Effect.gen(function* () {
const validationError = "Unknown model: anthropic/unknown";
const fake = yield* makeFakePi;
const { runtime } = yield* openRuntime(fake);
const providerThread = yield* runtime.ensureThread({
threadId: THREAD_ID,
modelSelection: modelSelection("default"),
runtimePolicy,
});
fake.deferNextModelSelection();
const starting = yield* startTurn(runtime, providerThread, "anthropic/unknown").pipe(
Effect.exit,
Effect.forkChild,
);
const pending = yield* fake.takeRequest("set_model");
yield* fake.emit({
type: "response",
id: pending.id,
command: "set_model",
success: false,
error: validationError,
});
const result = yield* Fiber.join(starting);
assert.equal(result._tag, "Failure");
if (Exit.isFailure(result)) assert.include(Cause.pretty(result.cause), validationError);
yield* startTurn(runtime, providerThread, "anthropic/model-b");
yield* fake.takeRequest("prompt");
assert.deepEqual(fake.promptModels, ["anthropic/model-b"]);
}).pipe(Effect.scoped, Effect.provide(layerTest)),
);

it.effect("retires an interrupted switch before accepting further requests", () =>
Effect.gen(function* () {
const fake = yield* makeFakePi;
Expand Down
136 changes: 73 additions & 63 deletions packages/provider-pi/src/server/adapter.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2548,81 +2548,91 @@ export const makePiAdapterV2 = Effect.fn("makePiAdapterV2")(function* (
return providerThread;
});

const applySelection = Effect.fnUntraced(function* (modelSelection: ModelSelection) {
const thinking = getModelSelectionStringOptionValue(modelSelection, "thinking");
if (modelSelection.model === PI_INHERIT_MODEL_SLUG) {
if (appliedModel !== null && baselineModel !== null) {
const restoredModel = yield* request({ type: "set_model", ...baselineModel });
const applySelection = Effect.fnUntraced(
function* (modelSelection: ModelSelection) {
const thinking = getModelSelectionStringOptionValue(modelSelection, "thinking");
if (modelSelection.model === PI_INHERIT_MODEL_SLUG) {
if (appliedModel !== null && baselineModel !== null) {
const restoredModel = yield* request({ type: "set_model", ...baselineModel });
appliedThinking = null;
contextWindow = rememberModelContextWindow(restoredModel);
appliedModel = null;
const updatedAt = yield* DateTime.now;
sessionEntity = { ...sessionEntity, model: PI_INHERIT_MODEL_SLUG, updatedAt };
yield* emit({
type: "provider_session.updated",
driver: PI_PROVIDER,
providerSession: sessionEntity,
});
}
if (
thinking === undefined &&
baselineThinking !== null &&
appliedThinking !== baselineThinking
) {
yield* request({ type: "set_thinking_level", level: baselineThinking });
appliedThinking = baselineThinking;
}
} else if (modelSelection.model !== appliedModel) {
const parsed = parsePiModelSlug(modelSelection.model);
if (parsed === null) {
return yield* protocolError(
`Pi model '${modelSelection.model}' must use provider/model format`,
);
}
const selectedModel = yield* request({
type: "set_model",
provider: parsed.provider,
modelId: parsed.modelId,
});
// Pi resets thinking to the target model's configured default on set_model.
appliedThinking = null;
contextWindow = rememberModelContextWindow(restoredModel);
appliedModel = null;
contextWindow = rememberModelContextWindow(selectedModel);
appliedModel = modelSelection.model;
const updatedAt = yield* DateTime.now;
sessionEntity = { ...sessionEntity, model: PI_INHERIT_MODEL_SLUG, updatedAt };
sessionEntity = { ...sessionEntity, model: modelSelection.model, updatedAt };
yield* emit({
type: "provider_session.updated",
driver: PI_PROVIDER,
providerSession: sessionEntity,
});
}
if (
thinking === undefined &&
baselineThinking !== null &&
appliedThinking !== baselineThinking
thinking !== undefined &&
thinking !== appliedThinking &&
PI_THINKING_LEVELS.has(thinking)
) {
yield* request({ type: "set_thinking_level", level: baselineThinking });
appliedThinking = baselineThinking;
}
} else if (modelSelection.model !== appliedModel) {
const parsed = parsePiModelSlug(modelSelection.model);
if (parsed === null) {
return yield* protocolError(
`Pi model '${modelSelection.model}' must use provider/model format`,
);
yield* request({ type: "set_thinking_level", level: thinking });
appliedThinking = thinking;
}
const selectedModel = yield* request({
type: "set_model",
provider: parsed.provider,
modelId: parsed.modelId,
});
// Pi resets thinking to the target model's configured default on set_model.
appliedThinking = null;
contextWindow = rememberModelContextWindow(selectedModel);
appliedModel = modelSelection.model;
const updatedAt = yield* DateTime.now;
sessionEntity = { ...sessionEntity, model: modelSelection.model, updatedAt };
yield* emit({
type: "provider_session.updated",
driver: PI_PROVIDER,
providerSession: sessionEntity,
});
}
if (
thinking !== undefined &&
thinking !== appliedThinking &&
PI_THINKING_LEVELS.has(thinking)
) {
yield* request({ type: "set_thinking_level", level: thinking });
appliedThinking = thinking;
}
const requestedTier = getModelSelectionStringOptionValue(modelSelection, "serviceTier");
const serviceTier =
requestedTier === undefined || requestedTier === PI_STANDARD_SERVICE_TIER
? null
: requestedTier;
if (serviceTier !== null || serviceTierMayBeSet) {
serviceTierMayBeSet = true;
const result = yield* request({
type: "prompt",
message: `/${PI_T3_SERVICE_TIER_COMMAND} ${serviceTier ?? ""}`.trimEnd(),
});
// Pi before 0.99 acknowledges a handled extension command without a disposition.
const disposition = recordString(result, "disposition");
if (disposition !== undefined && disposition !== "handled") {
return yield* protocolError("T3 Code's Pi extension did not apply the service tier");
const requestedTier = getModelSelectionStringOptionValue(modelSelection, "serviceTier");
const serviceTier =
requestedTier === undefined || requestedTier === PI_STANDARD_SERVICE_TIER
? null
: requestedTier;
if (serviceTier !== null || serviceTierMayBeSet) {
serviceTierMayBeSet = true;
const result = yield* request({
type: "prompt",
message: `/${PI_T3_SERVICE_TIER_COMMAND} ${serviceTier ?? ""}`.trimEnd(),
});
// Pi before 0.99 acknowledges a handled extension command without a disposition.
const disposition = recordString(result, "disposition");
if (disposition !== undefined && disposition !== "handled") {
return yield* protocolError("T3 Code's Pi extension did not apply the service tier");
}
serviceTierMayBeSet = serviceTier !== null;
}
serviceTierMayBeSet = serviceTier !== null;
}
});
},
(effect) =>
effect.pipe(
// A timed-out or interrupted selection may still mutate Pi after the caller stops waiting.
Effect.tapError((error) =>
error._tag === "PiRpcTimeoutError" ? connection.terminate : Effect.void,
),
Effect.onInterrupt(() => connection.terminate),
),
);

const resolvePromptPayload = Effect.fnUntraced(function* (
text: string,
Expand Down
Loading