diff --git a/src/agents/embedded-agent-runner/run.overflow-compaction.loop.test.ts b/src/agents/embedded-agent-runner/run.overflow-compaction.loop.test.ts index 37f43530521e..61d63d05d2c9 100644 --- a/src/agents/embedded-agent-runner/run.overflow-compaction.loop.test.ts +++ b/src/agents/embedded-agent-runner/run.overflow-compaction.loop.test.ts @@ -14,6 +14,7 @@ import { mockedIsCompactionFailureError, mockedIsLikelyContextOverflowError, mockedLog, + mockedMarkAuthProfileSuccess, mockedResolveModelAsync, mockedRunEmbeddedAttempt, mockedSessionLikelyHasOversizedToolResults, @@ -187,6 +188,22 @@ describe("overflow compaction in run loop", () => { expect(requireMockCallArg(mockedRunEmbeddedAttempt, 0).thinkLevel).toBe("adaptive"); }); + it("does not wait for post-run auth-profile success bookkeeping before returning", async () => { + let resolveSuccess!: () => void; + const successPromise = new Promise((resolve) => { + resolveSuccess = resolve; + }); + mockedMarkAuthProfileSuccess.mockReturnValueOnce(successPromise); + mockedRunEmbeddedAttempt.mockResolvedValueOnce(makeAttemptResult()); + + const result = await runEmbeddedAgent(baseParams); + + expect(result.meta.error).toBeUndefined(); + expect(mockedMarkAuthProfileSuccess).toHaveBeenCalledTimes(1); + resolveSuccess(); + await successPromise; + }); + it("continues from transcript after compaction when the current inbound message was persisted", async () => { const overflowError = makeOverflowError(); diff --git a/src/agents/embedded-agent-runner/run.ts b/src/agents/embedded-agent-runner/run.ts index 9e7e95905ff1..2b718aa0c38c 100644 --- a/src/agents/embedded-agent-runner/run.ts +++ b/src/agents/embedded-agent-runner/run.ts @@ -31,6 +31,7 @@ import { import { sleepWithAbort } from "../../infra/backoff.js"; import { freezeDiagnosticTraceContext } from "../../infra/diagnostic-trace-context.js"; import { formatErrorMessage, toErrorObject } from "../../infra/errors.js"; +import { redactIdentifier } from "../../logging/redact-identifier.js"; import { buildAgentHookContextChannelFields } from "../../plugins/hook-agent-context.js"; import { getGlobalHookRunner } from "../../plugins/hook-runner-global.js"; import { resolveProviderAuthProfileId } from "../../plugins/provider-runtime.js"; @@ -599,6 +600,8 @@ function resolveInitialEmbeddedRunModel(params: { }; } +const POST_RUN_AUTH_PROFILE_SUCCESS_SLOW_MS = 1_000; + export function runEmbeddedAgent( paramsInput: RunEmbeddedAgentParams, ): Promise { @@ -1733,6 +1736,48 @@ async function runEmbeddedAgentInternal( modelId: failure.modelId, }); }; + const markAuthProfileSuccessAfterRun = () => { + if (!lastProfileId) { + return; + } + const successProfileId = lastProfileId; + const safeSuccessProfileId = redactIdentifier(successProfileId, { len: 12 }); + const successProvider = resolveAuthProfileStateProvider( + profileFailureStore, + successProfileId, + provider, + ); + const successStarted = Date.now(); + void markAuthProfileSuccess({ + store: profileFailureStore, + provider: successProvider, + profileId: successProfileId, + agentDir: params.agentDir, + }) + .then(() => { + const durationMs = Date.now() - successStarted; + if (durationMs >= POST_RUN_AUTH_PROFILE_SUCCESS_SLOW_MS) { + log.warn( + `post-run auth-profile success bookkeeping completed after ${durationMs}ms: ` + + `runId=${params.runId} sessionId=${params.sessionId} ` + + `provider=${sanitizeForLog(successProvider)} profileId=${safeSuccessProfileId}`, + ); + } else if (log.isEnabled("trace")) { + log.trace( + `post-run auth-profile success bookkeeping completed: ` + + `runId=${params.runId} sessionId=${params.sessionId} durationMs=${durationMs}`, + ); + } + }) + .catch((err: unknown) => { + log.warn( + `post-run auth-profile success bookkeeping failed: ` + + `runId=${params.runId} sessionId=${params.sessionId} ` + + `provider=${sanitizeForLog(successProvider)} profileId=${safeSuccessProfileId} ` + + `error=${formatErrorMessage(err)}`, + ); + }); + }; const resolveRunAuthProfileFailureReason = ( failoverReason: FailoverReason | null, opts?: { providerStarted?: boolean; transientRateLimit?: boolean }, @@ -4036,18 +4081,7 @@ async function runEmbeddedAgentInternal( log.debug( `embedded run done: runId=${params.runId} sessionId=${params.sessionId} durationMs=${Date.now() - started} aborted=${aborted}`, ); - if (lastProfileId) { - await markAuthProfileSuccess({ - store: profileFailureStore, - provider: resolveAuthProfileStateProvider( - profileFailureStore, - lastProfileId, - provider, - ), - profileId: lastProfileId, - agentDir: params.agentDir, - }); - } + markAuthProfileSuccessAfterRun(); const replayInvalid = resolveReplayInvalidForAttempt(null); const livenessState = attempt.yieldDetected ? "paused"