From 848ffe90e6551a7726ecfc10d4eaf4a18837efb5 Mon Sep 17 00:00:00 2001
From: Peter Steinberger <steipete@gmail.com>
Date: Fri, 8 May 2026 20:02:43 +0100
Subject: [PATCH] test: tighten plugin contract async callbacks

---
 src/agents/model-auth.profiles.test.ts            |  6 ++----
 src/plugins/contracts/host-hooks.contract.test.ts | 10 ++++++++--
 2 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/src/agents/model-auth.profiles.test.ts b/src/agents/model-auth.profiles.test.ts
index 1d35c9cd0f9..0f9a1b46cdc 100644
--- a/src/agents/model-auth.profiles.test.ts
+++ b/src/agents/model-auth.profiles.test.ts
@@ -372,15 +372,13 @@ async function resolveDemoLocalApiKey(params: {
   storedKeys: string[];
   configuredApiKey: string;
 }) {
-  let resolved!: Awaited<ReturnType<typeof resolveApiKeyForProvider>>;
-  await withEnvAsync({ DEMO_LOCAL_API_KEY: params.envApiKey }, async () => {
-    resolved = await resolveApiKeyForProvider({
+  return await withEnvAsync({ DEMO_LOCAL_API_KEY: params.envApiKey }, async () => {
+    return await resolveApiKeyForProvider({
       provider: "demo-local",
       store: buildDemoLocalStore(params.storedKeys),
       cfg: buildDemoLocalProviderCfg(params.configuredApiKey),
     });
   });
-  return resolved;
 }
 
 describe("getApiKeyForModel", () => {
diff --git a/src/plugins/contracts/host-hooks.contract.test.ts b/src/plugins/contracts/host-hooks.contract.test.ts
index cbb6e60693d..4b5cbdb21f3 100644
--- a/src/plugins/contracts/host-hooks.contract.test.ts
+++ b/src/plugins/contracts/host-hooks.contract.test.ts
@@ -1976,7 +1976,7 @@ describe("host-hook fixture plugin contract", () => {
 
   it("does not let stale scheduler cleanup delete a newer job generation", async () => {
     let releaseCleanup: (() => void) | undefined;
-    let markCleanupStarted!: () => void;
+    let markCleanupStarted: (() => void) | undefined;
     const cleanupStartedPromise = new Promise<void>((resolve) => {
       markCleanupStarted = resolve;
     });
@@ -1994,6 +1994,9 @@ describe("host-hook fixture plugin contract", () => {
           sessionKey: "agent:main:main",
           kind: "monitor",
           cleanup: async () => {
+            if (!markCleanupStarted) {
+              throw new Error("Expected scheduler cleanup start callback to be initialized");
+            }
             markCleanupStarted();
             await new Promise<void>((resolve) => {
               releaseCleanup = resolve;
@@ -2027,7 +2030,10 @@ describe("host-hook fixture plugin contract", () => {
       },
     });
 
-    releaseCleanup?.();
+    if (!releaseCleanup) {
+      throw new Error("Expected scheduler cleanup release callback to be initialized");
+    }
+    releaseCleanup();
     await expect(cleanupPromise).resolves.toMatchObject({ failures: [] });
     expect(listPluginSessionSchedulerJobs("scheduler-race")).toEqual([
       {