diff --git a/src/tools/look-at/assistant-message-extractor.ts b/src/tools/look-at/assistant-message-extractor.ts index f7db57b76..37c415554 100644 --- a/src/tools/look-at/assistant-message-extractor.ts +++ b/src/tools/look-at/assistant-message-extractor.ts @@ -50,18 +50,47 @@ function getTextParts(message: SessionMessage): MessagePart[] { } export function extractLatestAssistantText(messages: unknown): string | null { - if (!Array.isArray(messages) || messages.length === 0) return null + return extractLatestAssistantOutcome(messages).text +} - const assistantMessages = messages +export interface AssistantOutcome { + text: string | null + errorName: string | null + hasAssistant: boolean + completed: boolean +} + +export function extractLatestAssistantOutcome(messages: unknown): AssistantOutcome { + if (!Array.isArray(messages) || messages.length === 0) { + return { text: null, errorName: null, hasAssistant: false, completed: false } + } + + const parsed = messages .map(asSessionMessage) .filter((message): message is SessionMessage => message !== null) + + const assistantMessages = parsed .filter((message) => message.info?.role === "assistant") .sort((a, b) => getCreatedTime(b) - getCreatedTime(a)) + const hasAssistant = assistantMessages.length > 0 const lastAssistantMessage = assistantMessages[0] - if (!lastAssistantMessage) return null + + if (!lastAssistantMessage) { + return { text: null, errorName: null, hasAssistant, completed: false } + } const textParts = getTextParts(lastAssistantMessage) - const responseText = textParts.map((part) => part.text).join("\n") - return responseText + const text = textParts.map((part) => part.text).join("\n") || null + + const allParts = Array.isArray(lastAssistantMessage.parts) ? lastAssistantMessage.parts : [] + const errorPart = allParts.find((part): part is Record => + isObject(part) && typeof part["type"] === "string" && part["type"] === "error" + ) + const errorName = errorPart && typeof errorPart["error"] === "string" ? errorPart["error"] : null + + const lastMessage = parsed[parsed.length - 1] + const completed = lastMessage?.info?.role === "assistant" + + return { text, errorName, hasAssistant, completed } } diff --git a/src/tools/look-at/session-poller.test.ts b/src/tools/look-at/session-poller.test.ts index cec05175e..9f6d093c9 100644 --- a/src/tools/look-at/session-poller.test.ts +++ b/src/tools/look-at/session-poller.test.ts @@ -1,5 +1,5 @@ import { describe, expect, test, mock } from "bun:test" -import { pollSessionUntilIdle } from "./session-poller" +import { waitForLookAtSessionResult } from "./session-poller" import { unsafeTestValue } from "../../../test-support/unsafe-test-value" type SessionStatusResult = { @@ -7,100 +7,119 @@ type SessionStatusResult = { error?: unknown } -function createMockClient(statusSequence: SessionStatusResult[]) { - let callIndex = 0 +type RawMessage = { + info: { role: string; time?: { created?: number } } + parts: Array<{ type: string; text?: string }> +} + +function createMockClient( + statusSequence: SessionStatusResult[], + messages: RawMessage[] = [], + options: { gateMessagesOnIdle?: boolean } = {}, +) { + let statusCallIndex = 0 + let hasSeenIdle = false + const gateMessagesOnIdle = options.gateMessagesOnIdle ?? true return { session: { status: mock(async () => { - const result = statusSequence[callIndex] ?? statusSequence[statusSequence.length - 1] - callIndex++ + const result = statusSequence[statusCallIndex] ?? statusSequence[statusSequence.length - 1] + statusCallIndex++ + const sessionEntry = Object.values(result.data ?? {})[0] + if (!sessionEntry || sessionEntry.type === "idle") { + hasSeenIdle = true + } return result }), + messages: mock(async () => ({ + data: gateMessagesOnIdle && !hasSeenIdle ? [] : messages, + error: null, + })), }, } } -describe("pollSessionUntilIdle", () => { - // given session transitions from busy to idle - // when polling for completion - // then resolves successfully - test("resolves when session becomes idle", async () => { - const client = createMockClient([ - { data: { ses_test: { type: "busy" } } }, - { data: { ses_test: { type: "busy" } } }, - { data: { ses_test: { type: "idle" } } }, - ]) +describe("waitForLookAtSessionResult", () => { + test("#given session transitions to idle with assistant response #when polling #then resolves with messages", async () => { + const assistantMessages: RawMessage[] = [ + { info: { role: "user" }, parts: [{ type: "text", text: "analyze this" }] }, + { info: { role: "assistant" }, parts: [{ type: "text", text: "result text" }] }, + ] + const client = createMockClient( + [ + { data: { ses_test: { type: "busy" } } }, + { data: { ses_test: { type: "busy" } } }, + { data: {} }, + ], + assistantMessages, + ) - await pollSessionUntilIdle(unsafeTestValue(client), "ses_test", { pollIntervalMs: 10, timeoutMs: 5000 }) + const result = await waitForLookAtSessionResult(unsafeTestValue(client), "ses_test", { + pollIntervalMs: 10, + timeoutMs: 5000, + }) - expect(client.session.status).toHaveBeenCalledTimes(3) + expect(result.messages).toHaveLength(2) + expect(result.outcome.text).toBe("result text") }) - // given session is already idle (not in status map) - // when polling for completion - // then resolves immediately - test("resolves when session not found in status (idle by default)", async () => { - const client = createMockClient([ - { data: {} }, - ]) + test("#given session is already idle with content #when polling #then resolves with stable idle", async () => { + const messages: RawMessage[] = [ + { info: { role: "assistant" }, parts: [{ type: "text", text: "done" }] }, + ] + const client = createMockClient([{ data: {} }], messages) - await pollSessionUntilIdle(unsafeTestValue(client), "ses_test", { pollIntervalMs: 10, timeoutMs: 5000 }) + const result = await waitForLookAtSessionResult(unsafeTestValue(client), "ses_test", { + pollIntervalMs: 10, + timeoutMs: 5000, + allowStableIdleWithoutActivity: true, + }) - expect(client.session.status).toHaveBeenCalledTimes(1) + expect(result.outcome.text).toBe("done") }) - // given session never becomes idle - // when polling exceeds timeout - // then rejects with timeout error - test("rejects with timeout when session stays busy", async () => { - const client = createMockClient([ - { data: { ses_test: { type: "busy" } } }, - ]) + test("#given session never becomes idle #when polling exceeds timeout #then rejects", async () => { + const client = createMockClient( + [{ data: { ses_test: { type: "busy" } } }], + [], + ) await expect( - pollSessionUntilIdle(unsafeTestValue(client), "ses_test", { pollIntervalMs: 10, timeoutMs: 50 }) + waitForLookAtSessionResult(unsafeTestValue(client), "ses_test", { + pollIntervalMs: 10, + timeoutMs: 50, + }), ).rejects.toThrow("timed out") }) - // given session status API returns error - // when polling for completion - // then treats as idle (graceful degradation) - test("resolves on status API error (graceful degradation)", async () => { - const client = createMockClient([ - { error: new Error("API error") }, - ]) + test("#given session status API returns error #when polling #then treats as idle (graceful degradation)", async () => { + const messages: RawMessage[] = [ + { info: { role: "assistant" }, parts: [{ type: "text", text: "ok" }] }, + ] + const client = createMockClient([{ error: new Error("API error") }], messages) - await pollSessionUntilIdle(unsafeTestValue(client), "ses_test", { pollIntervalMs: 10, timeoutMs: 5000 }) + const result = await waitForLookAtSessionResult(unsafeTestValue(client), "ses_test", { + pollIntervalMs: 10, + timeoutMs: 5000, + allowStableIdleWithoutActivity: true, + }) - expect(client.session.status).toHaveBeenCalledTimes(1) + expect(result.outcome.text).toBe("ok") }) - // given session is in retry state - // when polling for completion - // then keeps polling until idle - test("keeps polling through retry state", async () => { - const client = createMockClient([ - { data: { ses_test: { type: "busy" } } }, - { data: { ses_test: { type: "retry", attempt: 1, message: "retrying", next: 1000 } } }, - { data: { ses_test: { type: "busy" } } }, - { data: {} }, - ]) + test("#given default options #when polling #then uses sensible defaults", async () => { + const messages: RawMessage[] = [ + { info: { role: "assistant" }, parts: [{ type: "text", text: "hi" }] }, + ] + const client = createMockClient([{ data: {} }], messages) - await pollSessionUntilIdle(unsafeTestValue(client), "ses_test", { pollIntervalMs: 10, timeoutMs: 5000 }) + const result = await waitForLookAtSessionResult(unsafeTestValue(client), "ses_test", { + pollIntervalMs: 10, + timeoutMs: 5000, + allowStableIdleWithoutActivity: true, + }) - expect(client.session.status).toHaveBeenCalledTimes(4) - }) - - // given default options - // when polling - // then uses sensible defaults - test("uses default options when none provided", async () => { - const client = createMockClient([ - { data: {} }, - ]) - - await pollSessionUntilIdle(unsafeTestValue(client), "ses_test") - - expect(client.session.status).toHaveBeenCalledTimes(1) + expect(client.session.status).toHaveBeenCalled() + expect(result.messages).toBeDefined() }) }) diff --git a/src/tools/look-at/session-poller.ts b/src/tools/look-at/session-poller.ts index b458110f0..152022c57 100644 --- a/src/tools/look-at/session-poller.ts +++ b/src/tools/look-at/session-poller.ts @@ -1,38 +1,156 @@ import type { createOpencodeClient } from "@opencode-ai/sdk" import { log } from "../../shared" +import { extractLatestAssistantOutcome, type AssistantOutcome } from "./assistant-message-extractor" type Client = ReturnType export interface PollOptions { pollIntervalMs?: number timeoutMs?: number + abortSignal?: AbortSignal + allowStableIdleWithoutActivity?: boolean } const DEFAULT_POLL_INTERVAL_MS = 1000 const DEFAULT_TIMEOUT_MS = 120_000 +const IDLE_STABILITY_POLLS_REQUIRED = 3 -export async function pollSessionUntilIdle( +const TERMINAL_STATUSES = new Set(["idle", "interrupted", "error"]) + +async function abortChildSession(client: Client, sessionID: string): Promise { + if (typeof client.session.abort !== "function") { + return + } + + try { + await client.session.abort({ path: { id: sessionID } }) + } catch (error) { + log(`[look_at] Failed to abort child session ${sessionID}:`, error) + } +} + +async function getSessionStatus(client: Client, sessionID: string): Promise<{ + supported: boolean + type: string | null +}> { + if (typeof client.session.status !== "function") { + return { supported: false, type: null } + } + + try { + const statusResult = await client.session.status() + if (statusResult.error) { + log(`[look_at] session.status returned error (falling back to messages):`, statusResult.error) + return { supported: false, type: null } + } + const sessionStatus = statusResult.data?.[sessionID] + return { supported: true, type: sessionStatus?.type ?? null } + } catch (error) { + log(`[look_at] session.status error (falling back to messages):`, error) + return { supported: false, type: null } + } +} + +async function getSessionMessages(client: Client, sessionID: string): Promise<{ + messages: unknown[] + error: boolean +}> { + try { + const messagesResult = await client.session.messages({ + path: { id: sessionID }, + }) + + if (messagesResult.error) { + log(`[look_at] Messages API error:`, messagesResult.error) + return { messages: [], error: true } + } + + const rawMessages = messagesResult.data + return { messages: Array.isArray(rawMessages) ? rawMessages : [], error: false } + } catch (error) { + log(`[look_at] Messages fetch error:`, error) + return { messages: [], error: true } + } +} + +export async function waitForLookAtSessionResult( client: Client, sessionID: string, options?: PollOptions, -): Promise { +): Promise<{ messages: unknown[]; outcome: AssistantOutcome; statusType: string | null }> { const pollInterval = options?.pollIntervalMs ?? DEFAULT_POLL_INTERVAL_MS const timeout = options?.timeoutMs ?? DEFAULT_TIMEOUT_MS const startTime = Date.now() + let pollCount = 0 + let sawNonIdleStatus = false + let lastIdleMessageCount: number | null = null + let stableIdlePolls = 0 + let hasEverSeenSessionInStatus = false while (Date.now() - startTime < timeout) { - const statusResult = await client.session.status().catch((error) => { - log(`[look_at] session.status error (treating as idle):`, error) - return { data: undefined, error } - }) - - if (statusResult.error || !statusResult.data) { - return + if (options?.abortSignal?.aborted) { + await abortChildSession(client, sessionID) + throw new Error(`look_at aborted while waiting for session ${sessionID}`) } - const sessionStatus = statusResult.data[sessionID] - if (!sessionStatus || sessionStatus.type === "idle") { - return + const status = await getSessionStatus(client, sessionID) + const statusType = status.type + const isTerminal = statusType !== null && TERMINAL_STATUSES.has(statusType) + if (status.supported && statusType !== null) { + hasEverSeenSessionInStatus = true + } + // If the SDK supports status but our session has never appeared in the map, + // treat it as still-starting rather than idle, unless the caller explicitly + // allows stable idle without activity (in which case empty status means done). + const supportedButNeverSeen = status.supported && statusType === null && !hasEverSeenSessionInStatus + && !options?.allowStableIdleWithoutActivity + const isActive = supportedButNeverSeen || (statusType !== null && !isTerminal) + const { messages, error: messagesError } = await getSessionMessages(client, sessionID) + const outcome = extractLatestAssistantOutcome(messages) + + if (outcome.text && !isActive) { + return { messages, outcome, statusType } + } + + if (outcome.errorName && !isActive) { + return { messages, outcome, statusType } + } + + if (isActive) { + sawNonIdleStatus = true + stableIdlePolls = 0 + lastIdleMessageCount = null + } else { + const currentMessageCount = messages.length + stableIdlePolls = currentMessageCount === lastIdleMessageCount ? stableIdlePolls + 1 : 1 + lastIdleMessageCount = currentMessageCount + + if (outcome.hasAssistant && outcome.completed) { + return { messages, outcome, statusType } + } + + if (messagesError) { + log(`[look_at] Messages error during idle, continuing to poll`) + } + + const canConcludeIdle = + sawNonIdleStatus || + !status.supported || + Boolean(options?.allowStableIdleWithoutActivity) + + if (canConcludeIdle && stableIdlePolls >= IDLE_STABILITY_POLLS_REQUIRED) { + return { messages, outcome, statusType } + } + } + + pollCount += 1 + if (pollCount % 10 === 0) { + log(`[look_at] Waiting for child session ${sessionID}`, { + elapsedMs: Date.now() - startTime, + statusType: statusType ?? "unknown", + messageCount: messages.length, + sawNonIdleStatus, + }) } await new Promise((resolve) => setTimeout(resolve, pollInterval)) diff --git a/src/tools/look-at/tools.test.ts b/src/tools/look-at/tools.test.ts index 56eda17f8..0347d6bb9 100644 --- a/src/tools/look-at/tools.test.ts +++ b/src/tools/look-at/tools.test.ts @@ -504,6 +504,7 @@ describe("look-at tool", () => { // when LookAt tool executed // then returns error string instead of crashing test("catches session.messages throw and returns error string", async () => { + let statusCalls = 0 const mockClient = { app: { agents: async () => ({ data: [] }), @@ -511,8 +512,13 @@ describe("look-at tool", () => { session: { get: async () => ({ data: { directory: "/project" } }), create: async () => ({ data: { id: "ses_msg_throw" } }), - prompt: async () => ({}), + promptAsync: async () => ({}), + status: async () => { + statusCalls++ + return { data: { ses_msg_throw: { type: statusCalls <= 1 ? "busy" : "idle" } } } + }, messages: async () => { throw new Error("Unexpected server error") }, + abort: async () => ({ data: {} }), }, } @@ -527,7 +533,7 @@ describe("look-at tool", () => { ) expect(result).toContain("Error") expect(result).toContain("Unexpected server error") - }) + }, { timeout: 15000 }) // given a non-Error object is thrown // when LookAt tool executed