From 46e02b94575b18f7c8a403bd538a83f05879488b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pe=C3=AFo=20Thibault?= Date: Sat, 7 Feb 2026 13:42:24 +0100 Subject: [PATCH] fix(hooks): switch session.prompt to promptAsync in all hooks --- src/hooks/atlas/index.ts | 2 +- src/hooks/ralph-loop/index.ts | 2 +- src/hooks/session-recovery/index.ts | 4 +- src/hooks/todo-continuation-enforcer.ts | 2 +- src/hooks/unstable-agent-babysitter/index.ts | 11 +- src/shared/model-suggestion-retry.test.ts | 161 +++++++++++++++++-- src/tools/call-omo-agent/tools.ts | 24 +-- src/tools/delegate-task/executor.ts | 28 ++-- 8 files changed, 190 insertions(+), 44 deletions(-) diff --git a/src/hooks/atlas/index.ts b/src/hooks/atlas/index.ts index e3ab910a..ffad0459 100644 --- a/src/hooks/atlas/index.ts +++ b/src/hooks/atlas/index.ts @@ -484,7 +484,7 @@ export function createAtlasHook( : undefined } - await ctx.client.session.prompt({ + await ctx.client.session.promptAsync({ path: { id: sessionID }, body: { agent: agent ?? "atlas", diff --git a/src/hooks/ralph-loop/index.ts b/src/hooks/ralph-loop/index.ts index c961fe66..3cc77edd 100644 --- a/src/hooks/ralph-loop/index.ts +++ b/src/hooks/ralph-loop/index.ts @@ -364,7 +364,7 @@ export function createRalphLoopHook( : undefined } - await ctx.client.session.prompt({ + await ctx.client.session.promptAsync({ path: { id: sessionID }, body: { ...(agent !== undefined ? { agent } : {}), diff --git a/src/hooks/session-recovery/index.ts b/src/hooks/session-recovery/index.ts index ffd09077..2aecee15 100644 --- a/src/hooks/session-recovery/index.ts +++ b/src/hooks/session-recovery/index.ts @@ -75,7 +75,7 @@ function extractResumeConfig(userMessage: MessageData | undefined, sessionID: st async function resumeSession(client: Client, config: ResumeConfig): Promise { try { - await client.session.prompt({ + await client.session.promptAsync({ path: { id: config.sessionID }, body: { parts: [{ type: "text", text: RECOVERY_RESUME_TEXT }], @@ -185,7 +185,7 @@ async function recoverToolResultMissing( })) try { - await client.session.prompt({ + await client.session.promptAsync({ path: { id: sessionID }, // @ts-expect-error - SDK types may not include tool_result parts body: { parts: toolResultParts }, diff --git a/src/hooks/todo-continuation-enforcer.ts b/src/hooks/todo-continuation-enforcer.ts index 35e1df9d..3e3736be 100644 --- a/src/hooks/todo-continuation-enforcer.ts +++ b/src/hooks/todo-continuation-enforcer.ts @@ -245,7 +245,7 @@ ${todoList}` try { log(`[${HOOK_NAME}] Injecting continuation`, { sessionID, agent: agentName, model, incompleteCount: freshIncompleteCount }) - await ctx.client.session.prompt({ + await ctx.client.session.promptAsync({ path: { id: sessionID }, body: { agent: agentName, diff --git a/src/hooks/unstable-agent-babysitter/index.ts b/src/hooks/unstable-agent-babysitter/index.ts index 1a8b4057..e492281a 100644 --- a/src/hooks/unstable-agent-babysitter/index.ts +++ b/src/hooks/unstable-agent-babysitter/index.ts @@ -25,6 +25,15 @@ type BabysitterContext = { } query?: { directory?: string } }) => Promise + promptAsync: (args: { + path: { id: string } + body: { + parts: Array<{ type: "text"; text: string }> + agent?: string + model?: { providerID: string; modelID: string } + } + query?: { directory?: string } + }) => Promise } } } @@ -218,7 +227,7 @@ export function createUnstableAgentBabysitterHook(ctx: BabysitterContext, option const { agent, model } = await resolveMainSessionTarget(ctx, mainSessionID) try { - await ctx.client.session.prompt({ + await ctx.client.session.promptAsync({ path: { id: mainSessionID }, body: { ...(agent ? { agent } : {}), diff --git a/src/shared/model-suggestion-retry.test.ts b/src/shared/model-suggestion-retry.test.ts index 7c7d40cc..93338bd1 100644 --- a/src/shared/model-suggestion-retry.test.ts +++ b/src/shared/model-suggestion-retry.test.ts @@ -212,9 +212,9 @@ describe("parseModelSuggestion", () => { describe("promptWithModelSuggestionRetry", () => { it("should succeed on first try without retry", async () => { - // given a client where prompt succeeds + // given a client where promptAsync succeeds const promptMock = mock(() => Promise.resolve()) - const client = { session: { prompt: promptMock } } + const client = { session: { promptAsync: promptMock } } // when calling promptWithModelSuggestionRetry await promptWithModelSuggestionRetry(client as any, { @@ -225,21 +225,158 @@ describe("promptWithModelSuggestionRetry", () => { }, }) - // then should call prompt exactly once + // then should call promptAsync exactly once expect(promptMock).toHaveBeenCalledTimes(1) }) - it("should retry with suggestion on model-not-found error", async () => { - // given a client that fails first with model-not-found, then succeeds - const promptMock = mock() - .mockRejectedValueOnce({ - name: "ProviderModelNotFoundError", - data: { - providerID: "anthropic", - modelID: "claude-sonet-4", - suggestions: ["claude-sonnet-4"], + it("should throw error from promptAsync directly on model-not-found error", async () => { + // given a client that fails with model-not-found error + const promptMock = mock().mockRejectedValueOnce({ + name: "ProviderModelNotFoundError", + data: { + providerID: "anthropic", + modelID: "claude-sonet-4", + suggestions: ["claude-sonnet-4"], + }, + }) + const client = { session: { promptAsync: promptMock } } + + // when calling promptWithModelSuggestionRetry + // then should throw the error without retrying + await expect( + promptWithModelSuggestionRetry(client as any, { + path: { id: "session-1" }, + body: { + agent: "explore", + parts: [{ type: "text", text: "hello" }], + model: { providerID: "anthropic", modelID: "claude-sonet-4" }, }, }) + ).rejects.toThrow() + + // and should call promptAsync only once + expect(promptMock).toHaveBeenCalledTimes(1) + }) + + it("should throw original error when no suggestion available", async () => { + // given a client that fails with a non-model-not-found error + const originalError = new Error("Connection refused") + const promptMock = mock().mockRejectedValueOnce(originalError) + const client = { session: { promptAsync: promptMock } } + + // when calling promptWithModelSuggestionRetry + // then should throw the original error + await expect( + promptWithModelSuggestionRetry(client as any, { + path: { id: "session-1" }, + body: { + parts: [{ type: "text", text: "hello" }], + model: { providerID: "anthropic", modelID: "claude-sonnet-4" }, + }, + }) + ).rejects.toThrow("Connection refused") + + expect(promptMock).toHaveBeenCalledTimes(1) + }) + + it("should throw error from promptAsync directly", async () => { + // given a client that fails with an error + const error = new Error("Still not found") + const promptMock = mock().mockRejectedValueOnce(error) + const client = { session: { promptAsync: promptMock } } + + // when calling promptWithModelSuggestionRetry + // then should throw the error + await expect( + promptWithModelSuggestionRetry(client as any, { + path: { id: "session-1" }, + body: { + parts: [{ type: "text", text: "hello" }], + model: { providerID: "anthropic", modelID: "claude-sonnet-4" }, + }, + }) + ).rejects.toThrow("Still not found") + + // and should call promptAsync only once + expect(promptMock).toHaveBeenCalledTimes(1) + }) + + it("should pass all body fields through to promptAsync", async () => { + // given a client where promptAsync succeeds + const promptMock = mock().mockResolvedValueOnce(undefined) + const client = { session: { promptAsync: promptMock } } + + // when calling with additional body fields + await promptWithModelSuggestionRetry(client as any, { + path: { id: "session-1" }, + body: { + agent: "explore", + system: "You are a helpful agent", + tools: { task: false }, + parts: [{ type: "text", text: "hello" }], + model: { providerID: "anthropic", modelID: "claude-sonnet-4" }, + variant: "max", + }, + }) + + // then call should pass all fields through unchanged + const call = promptMock.mock.calls[0][0] + expect(call.body.agent).toBe("explore") + expect(call.body.system).toBe("You are a helpful agent") + expect(call.body.tools).toEqual({ task: false }) + expect(call.body.variant).toBe("max") + expect(call.body.model).toEqual({ + providerID: "anthropic", + modelID: "claude-sonnet-4", + }) + }) + + it("should throw string error message from promptAsync", async () => { + // given a client that fails with a string error + const promptMock = mock().mockRejectedValueOnce( + new Error("Model not found: anthropic/claude-sonet-4. Did you mean: claude-sonnet-4?") + ) + const client = { session: { promptAsync: promptMock } } + + // when calling promptWithModelSuggestionRetry + // then should throw the error + await expect( + promptWithModelSuggestionRetry(client as any, { + path: { id: "session-1" }, + body: { + parts: [{ type: "text", text: "hello" }], + model: { providerID: "anthropic", modelID: "claude-sonnet-4" }, + }, + }) + ).rejects.toThrow() + + // and should call promptAsync only once + expect(promptMock).toHaveBeenCalledTimes(1) + }) + + it("should throw error when no model in original request", async () => { + // given a client that fails with an error + const modelNotFoundError = new Error( + "Model not found: anthropic/claude-sonet-4. Did you mean: claude-sonnet-4?" + ) + const promptMock = mock().mockRejectedValueOnce(modelNotFoundError) + const client = { session: { promptAsync: promptMock } } + + // when calling without model in body + // then should throw the error + await expect( + promptWithModelSuggestionRetry(client as any, { + path: { id: "session-1" }, + body: { + parts: [{ type: "text", text: "hello" }], + }, + }) + ).rejects.toThrow() + + // and should call promptAsync only once + expect(promptMock).toHaveBeenCalledTimes(1) + }) +}) .mockResolvedValueOnce(undefined) const client = { session: { prompt: promptMock } } diff --git a/src/tools/call-omo-agent/tools.ts b/src/tools/call-omo-agent/tools.ts index f18a6eec..02c88052 100644 --- a/src/tools/call-omo-agent/tools.ts +++ b/src/tools/call-omo-agent/tools.ts @@ -219,18 +219,18 @@ Original error: ${createResult.error}` log(`[call_omo_agent] Sending prompt to session ${sessionID}`) log(`[call_omo_agent] Prompt text:`, args.prompt.substring(0, 100)) - try { - await ctx.client.session.prompt({ - path: { id: sessionID }, - body: { - agent: args.subagent_type, - tools: { - ...getAgentToolRestrictions(args.subagent_type), - task: false, - }, - parts: [{ type: "text", text: args.prompt }], - }, - }) + try { + await ctx.client.session.promptAsync({ + path: { id: sessionID }, + body: { + agent: args.subagent_type, + tools: { + ...getAgentToolRestrictions(args.subagent_type), + task: false, + }, + parts: [{ type: "text", text: args.prompt }], + }, + }) } catch (error) { const errorMessage = error instanceof Error ? error.message : String(error) log(`[call_omo_agent] Prompt error:`, errorMessage) diff --git a/src/tools/delegate-task/executor.ts b/src/tools/delegate-task/executor.ts index bc6fd170..15d9983f 100644 --- a/src/tools/delegate-task/executor.ts +++ b/src/tools/delegate-task/executor.ts @@ -211,20 +211,20 @@ export async function executeSyncContinuation( : undefined } - await client.session.prompt({ - path: { id: args.session_id! }, - body: { - ...(resumeAgent !== undefined ? { agent: resumeAgent } : {}), - ...(resumeModel !== undefined ? { model: resumeModel } : {}), - tools: { - ...(resumeAgent ? getAgentToolRestrictions(resumeAgent) : {}), - task: false, - call_omo_agent: true, - question: false, - }, - parts: [{ type: "text", text: args.prompt }], - }, - }) + await client.session.promptAsync({ + path: { id: args.session_id! }, + body: { + ...(resumeAgent !== undefined ? { agent: resumeAgent } : {}), + ...(resumeModel !== undefined ? { model: resumeModel } : {}), + tools: { + ...(resumeAgent ? getAgentToolRestrictions(resumeAgent) : {}), + task: false, + call_omo_agent: true, + question: false, + }, + parts: [{ type: "text", text: args.prompt }], + }, + }) } catch (promptError) { if (toastManager) { toastManager.removeTask(taskId)