fix(agents): honor explicit rate-limit cooldown probes in fallback runs

This commit is contained in:
Vignesh Natarajan
2026-03-05 20:02:36 -08:00
parent ce71fac7d6
commit d45353f95b
14 changed files with 150 additions and 25 deletions

View File

@@ -186,7 +186,7 @@ export async function runAgentTurnWithFallback(params: {
const onToolResult = params.opts?.onToolResult;
const fallbackResult = await runWithModelFallback({
...resolveModelFallbackOptions(params.followupRun.run),
run: (provider, model) => {
run: (provider, model, runOptions) => {
// Notify that model selection is complete (including after fallback).
// This allows responsePrefix template interpolation with the actual model.
params.opts?.onModelSelected?.({
@@ -304,6 +304,7 @@ export async function runAgentTurnWithFallback(params: {
model,
runId,
authProfile,
allowRateLimitCooldownProbe: runOptions?.allowRateLimitCooldownProbe,
});
return (async () => {
const result = await runEmbeddedPiAgent({

View File

@@ -474,7 +474,7 @@ export async function runMemoryFlushIfNeeded(params: {
try {
await runWithModelFallback({
...resolveModelFallbackOptions(params.followupRun.run),
run: async (provider, model) => {
run: async (provider, model, runOptions) => {
const { authProfile, embeddedContext, senderContext } = buildEmbeddedRunContexts({
run: params.followupRun.run,
sessionCtx: params.sessionCtx,
@@ -487,6 +487,7 @@ export async function runMemoryFlushIfNeeded(params: {
model,
runId: flushRunId,
authProfile,
allowRateLimitCooldownProbe: runOptions?.allowRateLimitCooldownProbe,
});
const result = await runEmbeddedPiAgent({
...embeddedContext,

View File

@@ -166,6 +166,7 @@ export function buildEmbeddedRunBaseParams(params: {
model: string;
runId: string;
authProfile: ReturnType<typeof resolveProviderScopedAuthProfile>;
allowRateLimitCooldownProbe?: boolean;
}) {
return {
sessionFile: params.run.sessionFile,
@@ -186,6 +187,7 @@ export function buildEmbeddedRunBaseParams(params: {
bashElevated: params.run.bashElevated,
timeoutMs: params.run.timeoutMs,
runId: params.runId,
allowRateLimitCooldownProbe: params.allowRateLimitCooldownProbe,
};
}

View File

@@ -157,7 +157,7 @@ export function createFollowupRunner(params: {
agentId: queued.run.agentId,
sessionKey: queued.run.sessionKey,
}),
run: async (provider, model) => {
run: async (provider, model, runOptions) => {
const authProfile = resolveRunAuthProfile(queued.run, provider);
const result = await runEmbeddedPiAgent({
sessionId: queued.run.sessionId,
@@ -200,6 +200,7 @@ export function createFollowupRunner(params: {
bashElevated: queued.run.bashElevated,
timeoutMs: queued.run.timeoutMs,
runId,
allowRateLimitCooldownProbe: runOptions?.allowRateLimitCooldownProbe,
blockReplyBreak: queued.run.blockReplyBreak,
bootstrapPromptWarningSignaturesSeen,
bootstrapPromptWarningSignature: