Upgrade to Ink 6 and React 19 (#2096)

Co-authored-by: jacob314 <jacob314@gmail.com>
This commit is contained in:
Sandy Tao
2025-06-27 16:39:54 -07:00
committed by GitHub
parent 19d2a0fb35
commit 150df382f8
18 changed files with 1129 additions and 1571 deletions

View File

@@ -351,18 +351,6 @@ const App = ({ config, settings, startupWarnings = [] }: AppProps) => {
// the user starts interacting with the app.
enteringConstrainHeightMode = true;
setConstrainHeight(true);
// If our pending history item happens to exceed the terminal height we will most likely need to refresh
// our static collection to ensure no duplication or tearing. This is currently working around a core bug
// in Ink which we have a PR out to fix: https://github.com/vadimdemedes/ink/pull/717
if (pendingHistoryItemRef.current && pendingHistoryItems.length > 0) {
const pendingItemDimensions = measureElement(
pendingHistoryItemRef.current,
);
if (pendingItemDimensions.height > availableTerminalHeight) {
refreshStatic();
}
}
}
if (key.ctrl && input === 'o') {
@@ -530,23 +518,6 @@ const App = ({ config, settings, startupWarnings = [] }: AppProps) => {
};
}, [terminalWidth, terminalHeight, refreshStatic]);
useEffect(() => {
if (!pendingHistoryItems.length) {
return;
}
const pendingItemDimensions = measureElement(
pendingHistoryItemRef.current!,
);
// If our pending history item happens to exceed the terminal height we will most likely need to refresh
// our static collection to ensure no duplication or tearing. This is currently working around a core bug
// in Ink which we have a PR out to fix: https://github.com/vadimdemedes/ink/pull/717
if (pendingItemDimensions.height > availableTerminalHeight) {
setStaticNeedsRefresh(true);
}
}, [pendingHistoryItems.length, availableTerminalHeight, streamingState]);
useEffect(() => {
if (streamingState === StreamingState.Idle && staticNeedsRefresh) {
setStaticNeedsRefresh(false);

View File

@@ -116,13 +116,15 @@ export const MaxSizedBox: React.FC<MaxSizedBoxProps> = ({
throw new Error('maxWidth must be defined when maxHeight is set.');
}
function visitRows(element: React.ReactNode) {
if (!React.isValidElement(element)) {
if (!React.isValidElement<{ children?: React.ReactNode }>(element)) {
return;
}
if (element.type === Fragment) {
React.Children.forEach(element.props.children, visitRows);
return;
}
if (element.type === Box) {
layoutInkElementAsStyledText(element, maxWidth!, laidOutStyledText);
return;
@@ -246,7 +248,10 @@ interface Row {
* @returns An array of `Row` objects.
*/
function visitBoxRow(element: React.ReactNode): Row {
if (!React.isValidElement(element) || element.type !== Box) {
if (
!React.isValidElement<{ children?: React.ReactNode }>(element) ||
element.type !== Box
) {
debugReportError(
`All children of MaxSizedBox must be <Box> elements`,
element,
@@ -258,7 +263,15 @@ function visitBoxRow(element: React.ReactNode): Row {
}
if (enableDebugLog) {
const boxProps = element.props;
const boxProps = element.props as {
children?: React.ReactNode | undefined;
readonly flexDirection?:
| 'row'
| 'column'
| 'row-reverse'
| 'column-reverse'
| undefined;
};
// Ensure the Box has no props other than the default ones and key.
let maxExpectedProps = 4;
if (boxProps.children !== undefined) {
@@ -323,14 +336,13 @@ function visitBoxRow(element: React.ReactNode): Row {
return;
}
if (!React.isValidElement(element)) {
if (!React.isValidElement<{ children?: React.ReactNode }>(element)) {
debugReportError('Invalid element.', element);
return;
}
if (element.type === Fragment) {
const fragmentChildren = element.props.children;
React.Children.forEach(fragmentChildren, (child) =>
React.Children.forEach(element.props.children, (child) =>
visitRowChild(child, parentProps),
);
return;

View File

@@ -6,6 +6,7 @@
import { type MutableRefObject } from 'react';
import { render } from 'ink-testing-library';
import { renderHook } from '@testing-library/react';
import { act } from 'react-dom/test-utils';
import { SessionStatsProvider, useSessionStats } from './SessionContext.js';
import { describe, it, expect, vi } from 'vitest';
@@ -223,21 +224,16 @@ describe('SessionStatsContext', () => {
});
it('should throw an error when useSessionStats is used outside of a provider', () => {
// Suppress the expected console error during this test.
const errorSpy = vi.spyOn(console, 'error').mockImplementation(() => {});
// Suppress console.error for this test since we expect an error
const consoleSpy = vi.spyOn(console, 'error').mockImplementation(() => {});
const contextRef = { current: undefined };
// We expect rendering to fail, which React will catch and log as an error.
render(<TestHarness contextRef={contextRef} />);
// Assert that the first argument of the first call to console.error
// contains the expected message. This is more robust than checking
// the exact arguments, which can be affected by React/JSDOM internals.
expect(errorSpy.mock.calls[0][0]).toContain(
'useSessionStats must be used within a SessionStatsProvider',
);
errorSpy.mockRestore();
try {
// Expect renderHook itself to throw when the hook is used outside a provider
expect(() => {
renderHook(() => useSessionStats());
}).toThrow('useSessionStats must be used within a SessionStatsProvider');
} finally {
consoleSpy.mockRestore();
}
});
});

View File

@@ -25,9 +25,12 @@ export function useConsoleMessages(): UseConsoleMessagesReturn {
return;
}
const newMessagesToAdd = messageQueueRef.current;
messageQueueRef.current = [];
setConsoleMessages((prevMessages) => {
const newMessages = [...prevMessages];
messageQueueRef.current.forEach((queuedMessage) => {
newMessagesToAdd.forEach((queuedMessage) => {
if (
newMessages.length > 0 &&
newMessages[newMessages.length - 1].type === queuedMessage.type &&
@@ -42,7 +45,6 @@ export function useConsoleMessages(): UseConsoleMessagesReturn {
return newMessages;
});
messageQueueRef.current = [];
messageQueueTimeoutRef.current = null; // Allow next scheduling
}, []);

View File

@@ -496,13 +496,17 @@ describe('useGeminiStream', () => {
} as TrackedCompletedToolCall, // Treat error as a form of completion for submission
];
// 1. On the first render, there are no tool calls.
mockUseReactToolScheduler.mockReturnValue([
[],
mockScheduleToolCalls,
mockMarkToolsAsSubmitted,
]);
const { rerender } = renderHook(() =>
// Capture the onComplete callback
let capturedOnComplete:
| ((completedTools: TrackedToolCall[]) => Promise<void>)
| null = null;
mockUseReactToolScheduler.mockImplementation((onComplete) => {
capturedOnComplete = onComplete;
return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted];
});
renderHook(() =>
useGeminiStream(
new MockedGeminiClientClass(mockConfig),
[],
@@ -518,16 +522,11 @@ describe('useGeminiStream', () => {
),
);
// 2. Before the second render, change the mock to return the completed tools.
mockUseReactToolScheduler.mockReturnValue([
completedToolCalls,
mockScheduleToolCalls,
mockMarkToolsAsSubmitted,
]);
// 3. Trigger a re-render. The hook will now receive the completed tools, causing the effect to run.
act(() => {
rerender();
// Trigger the onComplete callback with completed tools
await act(async () => {
if (capturedOnComplete) {
await capturedOnComplete(completedToolCalls);
}
});
await waitFor(() => {
@@ -561,13 +560,17 @@ describe('useGeminiStream', () => {
];
const client = new MockedGeminiClientClass(mockConfig);
// 1. First render: no tool calls.
mockUseReactToolScheduler.mockReturnValue([
[],
mockScheduleToolCalls,
mockMarkToolsAsSubmitted,
]);
const { rerender } = renderHook(() =>
// Capture the onComplete callback
let capturedOnComplete:
| ((completedTools: TrackedToolCall[]) => Promise<void>)
| null = null;
mockUseReactToolScheduler.mockImplementation((onComplete) => {
capturedOnComplete = onComplete;
return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted];
});
renderHook(() =>
useGeminiStream(
client,
[],
@@ -583,16 +586,11 @@ describe('useGeminiStream', () => {
),
);
// 2. Second render: tool calls are now cancelled.
mockUseReactToolScheduler.mockReturnValue([
cancelledToolCalls,
mockScheduleToolCalls,
mockMarkToolsAsSubmitted,
]);
// 3. Trigger the re-render.
act(() => {
rerender();
// Trigger the onComplete callback with cancelled tools
await act(async () => {
if (capturedOnComplete) {
await capturedOnComplete(cancelledToolCalls);
}
});
await waitFor(() => {
@@ -685,7 +683,12 @@ describe('useGeminiStream', () => {
const initialToolCalls: TrackedToolCall[] = [
{
request: { callId: 'call1', name: 'tool1', args: {} },
request: {
callId: 'call1',
name: 'tool1',
args: {},
isClientInitiated: false,
},
status: 'executing',
responseSubmittedToGemini: false,
tool: {
@@ -711,36 +714,67 @@ describe('useGeminiStream', () => {
} as TrackedCompletedToolCall,
];
const { result, rerender, client } = renderTestHook(initialToolCalls);
// Capture the onComplete callback
let capturedOnComplete:
| ((completedTools: TrackedToolCall[]) => Promise<void>)
| null = null;
let currentToolCalls = initialToolCalls;
mockUseReactToolScheduler.mockImplementation((onComplete) => {
capturedOnComplete = onComplete;
return [
currentToolCalls,
mockScheduleToolCalls,
mockMarkToolsAsSubmitted,
];
});
const { result, rerender } = renderHook(() =>
useGeminiStream(
new MockedGeminiClientClass(mockConfig),
[],
mockAddItem,
mockSetShowHelp,
mockConfig,
mockOnDebugMessage,
mockHandleSlashCommand,
false,
() => 'vscode' as EditorType,
() => {},
() => Promise.resolve(),
),
);
// 1. Initial state should be Responding because a tool is executing.
expect(result.current.streamingState).toBe(StreamingState.Responding);
// 2. Rerender with the completed tool call.
// The useEffect should pick this up but hasn't called submitQuery yet.
// 2. Update the tool calls to completed state and rerender
currentToolCalls = completedToolCalls;
mockUseReactToolScheduler.mockImplementation((onComplete) => {
capturedOnComplete = onComplete;
return [
completedToolCalls,
mockScheduleToolCalls,
mockMarkToolsAsSubmitted,
];
});
act(() => {
rerender({
client,
history: [],
addItem: mockAddItem,
setShowHelp: mockSetShowHelp,
config: mockConfig,
onDebugMessage: mockOnDebugMessage,
handleSlashCommand:
mockHandleSlashCommand as unknown as typeof mockHandleSlashCommand,
shellModeActive: false,
loadedSettings: mockLoadedSettings,
// This is the key part of the test: update the toolCalls array
// to simulate the tool finishing.
toolCalls: completedToolCalls,
});
rerender();
});
// 3. The state should *still* be Responding, not Idle.
// This is because the completed tool's response has not been submitted yet.
expect(result.current.streamingState).toBe(StreamingState.Responding);
// 4. Wait for the useEffect to call submitQuery.
// 4. Trigger the onComplete callback to simulate tool completion
await act(async () => {
if (capturedOnComplete) {
await capturedOnComplete(completedToolCalls);
}
});
// 5. Wait for submitQuery to be called
await waitFor(() => {
expect(mockSendMessageStream).toHaveBeenCalledWith(
toolCallResponseParts,
@@ -748,7 +782,7 @@ describe('useGeminiStream', () => {
);
});
// 5. After submission, the state should remain Responding.
// 6. After submission, the state should remain Responding until the stream completes.
expect(result.current.streamingState).toBe(StreamingState.Responding);
});
@@ -929,14 +963,17 @@ describe('useGeminiStream', () => {
} as any,
};
// 1. Initial render state: no tool calls
mockUseReactToolScheduler.mockReturnValue([
[],
mockScheduleToolCalls,
mockMarkToolsAsSubmitted,
]);
// Capture the onComplete callback
let capturedOnComplete:
| ((completedTools: TrackedToolCall[]) => Promise<void>)
| null = null;
const { result, rerender } = renderHook(() =>
mockUseReactToolScheduler.mockImplementation((onComplete) => {
capturedOnComplete = onComplete;
return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted];
});
const { result } = renderHook(() =>
useGeminiStream(
new MockedGeminiClientClass(mockConfig),
[],
@@ -957,17 +994,11 @@ describe('useGeminiStream', () => {
await result.current.submitQuery('/memory add "test fact"');
});
// The command handler schedules the tool. Now we simulate the tool completing.
// 2. Before the next render, set the mock to return the completed tool.
mockUseReactToolScheduler.mockReturnValue([
[completedToolCall],
mockScheduleToolCalls,
mockMarkToolsAsSubmitted,
]);
// 3. Trigger a re-render to process the completed tool.
act(() => {
rerender();
// Trigger the onComplete callback with the completed client-initiated tool
await act(async () => {
if (capturedOnComplete) {
await capturedOnComplete([completedToolCall]);
}
});
// --- Assert the outcome ---
@@ -1007,13 +1038,17 @@ describe('useGeminiStream', () => {
} as any,
};
mockUseReactToolScheduler.mockReturnValue([
[completedToolCall],
mockScheduleToolCalls,
mockMarkToolsAsSubmitted,
]);
// Capture the onComplete callback
let capturedOnComplete:
| ((completedTools: TrackedToolCall[]) => Promise<void>)
| null = null;
const { rerender } = renderHook(() =>
mockUseReactToolScheduler.mockImplementation((onComplete) => {
capturedOnComplete = onComplete;
return [[], mockScheduleToolCalls, mockMarkToolsAsSubmitted];
});
renderHook(() =>
useGeminiStream(
new MockedGeminiClientClass(mockConfig),
[],
@@ -1029,8 +1064,11 @@ describe('useGeminiStream', () => {
),
);
act(() => {
rerender();
// Trigger the onComplete callback with the completed save_memory tool
await act(async () => {
if (capturedOnComplete) {
await capturedOnComplete([completedToolCall]);
}
});
await waitFor(() => {

View File

@@ -111,17 +111,21 @@ export const useGeminiStream = (
const [toolCalls, scheduleToolCalls, markToolsAsSubmitted] =
useReactToolScheduler(
(completedToolCallsFromScheduler) => {
async (completedToolCallsFromScheduler) => {
// This onComplete is called when ALL scheduled tools for a given batch are done.
if (completedToolCallsFromScheduler.length > 0) {
// Add the final state of these tools to the history for display.
// The new useEffect will handle submitting their responses.
addItem(
mapTrackedToolCallsToDisplay(
completedToolCallsFromScheduler as TrackedToolCall[],
),
Date.now(),
);
// Handle tool response submission immediately when tools complete
await handleCompletedTools(
completedToolCallsFromScheduler as TrackedToolCall[],
);
}
},
config,
@@ -570,40 +574,33 @@ export const useGeminiStream = (
],
);
/**
* Automatically submits responses for completed tool calls.
* This effect runs when `toolCalls` or `isResponding` changes.
* It ensures that tool responses are sent back to Gemini only when
* all processing for a given set of tools is finished and Gemini
* is not already generating a response.
*/
useEffect(() => {
const run = async () => {
const handleCompletedTools = useCallback(
async (completedToolCallsFromScheduler: TrackedToolCall[]) => {
if (isResponding) {
return;
}
const completedAndReadyToSubmitTools = toolCalls.filter(
(
tc: TrackedToolCall,
): tc is TrackedCompletedToolCall | TrackedCancelledToolCall => {
const isTerminalState =
tc.status === 'success' ||
tc.status === 'error' ||
tc.status === 'cancelled';
const completedAndReadyToSubmitTools =
completedToolCallsFromScheduler.filter(
(
tc: TrackedToolCall,
): tc is TrackedCompletedToolCall | TrackedCancelledToolCall => {
const isTerminalState =
tc.status === 'success' ||
tc.status === 'error' ||
tc.status === 'cancelled';
if (isTerminalState) {
const completedOrCancelledCall = tc as
| TrackedCompletedToolCall
| TrackedCancelledToolCall;
return (
!completedOrCancelledCall.responseSubmittedToGemini &&
completedOrCancelledCall.response?.responseParts !== undefined
);
}
return false;
},
);
if (isTerminalState) {
const completedOrCancelledCall = tc as
| TrackedCompletedToolCall
| TrackedCancelledToolCall;
return (
completedOrCancelledCall.response?.responseParts !== undefined
);
}
return false;
},
);
// Finalize any client-initiated tools as soon as they are done.
const clientTools = completedAndReadyToSubmitTools.filter(
@@ -630,15 +627,6 @@ export const useGeminiStream = (
);
}
// Only proceed with submitting to Gemini if ALL tools are complete.
const allToolsAreComplete =
toolCalls.length > 0 &&
toolCalls.length === completedAndReadyToSubmitTools.length;
if (!allToolsAreComplete) {
return;
}
const geminiTools = completedAndReadyToSubmitTools.filter(
(t) => !t.request.isClientInitiated,
);
@@ -693,17 +681,15 @@ export const useGeminiStream = (
submitQuery(mergePartListUnions(responsesToSend), {
isContinuation: true,
});
};
void run();
}, [
toolCalls,
isResponding,
submitQuery,
markToolsAsSubmitted,
addItem,
geminiClient,
performMemoryRefresh,
]);
},
[
isResponding,
submitQuery,
markToolsAsSubmitted,
geminiClient,
performMemoryRefresh,
],
);
const pendingHistoryItems = [
pendingHistoryItemRef.current,

View File

@@ -128,7 +128,7 @@ export function useReactToolScheduler(
}),
);
},
[],
[setToolCallsForDisplay],
);
const scheduler = useMemo(
@@ -152,7 +152,7 @@ export function useReactToolScheduler(
);
const schedule: ScheduleFn = useCallback(
async (
(
request: ToolCallRequestInfo | ToolCallRequestInfo[],
signal: AbortSignal,
) => {