Skip to content

Commit 4930a3d

Browse files
sestinjsestinj
committed
more chat tests
1 parent e609491 commit 4930a3d

File tree

2 files changed

+331
-3
lines changed

2 files changed

+331
-3
lines changed

‎gui/src/context/IdeMessenger.tsx‎

Lines changed: 277 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,277 @@
1+
import { ChatMessage, IDE, PromptLog } from "core";
2+
import type {
3+
FromWebviewProtocol,
4+
ToCoreProtocol,
5+
ToWebviewProtocol,
6+
} from "core/protocol";
7+
import { Message } from "core/protocol/messenger";
8+
import { MessageIde } from "core/protocol/messenger/messageIde";
9+
import {
10+
GeneratorReturnType,
11+
GeneratorYieldType,
12+
WebviewProtocolGeneratorMessage,
13+
WebviewSingleMessage,
14+
WebviewSingleProtocolMessage,
15+
} from "core/protocol/util";
16+
import { createContext } from "react";
17+
import { v4 as uuidv4 } from "uuid";
18+
import { isJetBrains } from "../util";
19+
20+
interface vscode {
21+
postMessage(message: any): vscode;
22+
}
23+
24+
declare const vscode: any;
25+
26+
export interface IIdeMessenger {
27+
post<T extends keyof FromWebviewProtocol>(
28+
messageType: T,
29+
data: FromWebviewProtocol[T][0],
30+
messageId?: string,
31+
attempt?: number,
32+
): void;
33+
34+
respond<T extends keyof ToWebviewProtocol>(
35+
messageType: T,
36+
data: ToWebviewProtocol[T][1],
37+
messageId: string,
38+
): void;
39+
40+
request<T extends keyof FromWebviewProtocol>(
41+
messageType: T,
42+
data: FromWebviewProtocol[T][0],
43+
): Promise<WebviewSingleProtocolMessage<T>>;
44+
45+
streamRequest<T extends keyof FromWebviewProtocol>(
46+
messageType: T,
47+
data: FromWebviewProtocol[T][0],
48+
cancelToken?: AbortSignal,
49+
): AsyncGenerator<
50+
GeneratorYieldType<FromWebviewProtocol[T][1]>[],
51+
GeneratorReturnType<FromWebviewProtocol[T][1]> | undefined
52+
>;
53+
54+
llmStreamChat(
55+
msg: ToCoreProtocol["llm/streamChat"][0],
56+
cancelToken: AbortSignal,
57+
): AsyncGenerator<ChatMessage[], PromptLog | undefined>;
58+
59+
ide: IDE;
60+
}
61+
62+
export class IdeMessenger implements IIdeMessenger {
63+
ide: IDE;
64+
65+
constructor() {
66+
this.ide = new MessageIde(
67+
async (messageType, data) => {
68+
const result = await this.request(messageType, data);
69+
if (result.status === "error") {
70+
throw new Error(result.error);
71+
}
72+
return result.content;
73+
},
74+
() => {},
75+
);
76+
}
77+
78+
private _postToIde(
79+
messageType: string,
80+
data: any,
81+
messageId: string = uuidv4(),
82+
) {
83+
if (typeof vscode === "undefined") {
84+
if (isJetBrains()) {
85+
if (window.postIntellijMessage === undefined) {
86+
console.log(
87+
"Unable to send message: postIntellijMessage is undefined. ",
88+
messageType,
89+
data,
90+
);
91+
throw new Error("postIntellijMessage is undefined");
92+
}
93+
window.postIntellijMessage?.(messageType, data, messageId);
94+
return;
95+
} else {
96+
console.log(
97+
"Unable to send message: vscode is undefined",
98+
messageType,
99+
data,
100+
);
101+
return;
102+
}
103+
}
104+
105+
const msg: Message = {
106+
messageId,
107+
messageType,
108+
data,
109+
};
110+
111+
vscode.postMessage(msg);
112+
}
113+
114+
post<T extends keyof FromWebviewProtocol>(
115+
messageType: T,
116+
data: FromWebviewProtocol[T][0],
117+
messageId?: string,
118+
attempt: number = 0,
119+
) {
120+
try {
121+
this._postToIde(messageType, data, messageId);
122+
} catch (error) {
123+
if (attempt < 5) {
124+
console.log(`Attempt ${attempt} failed. Retrying...`);
125+
setTimeout(
126+
() => this.post(messageType, data, messageId, attempt + 1),
127+
Math.pow(2, attempt) * 1000,
128+
);
129+
} else {
130+
console.error(
131+
"Max attempts reached. Message could not be sent.",
132+
error,
133+
);
134+
}
135+
}
136+
}
137+
138+
respond<T extends keyof ToWebviewProtocol>(
139+
messageType: T,
140+
data: ToWebviewProtocol[T][1],
141+
messageId: string,
142+
) {
143+
this._postToIde(messageType, data, messageId);
144+
}
145+
146+
request<T extends keyof FromWebviewProtocol>(
147+
messageType: T,
148+
data: FromWebviewProtocol[T][0],
149+
): Promise<WebviewSingleMessage<T>> {
150+
const messageId = uuidv4();
151+
152+
return new Promise((resolve) => {
153+
const handler = (event: any) => {
154+
if (event.data.messageId === messageId) {
155+
window.removeEventListener("message", handler);
156+
resolve(event.data.data as WebviewSingleMessage<T>);
157+
}
158+
};
159+
window.addEventListener("message", handler);
160+
161+
this.post(messageType, data, messageId);
162+
});
163+
}
164+
165+
/**
166+
* Because of weird type stuff, we're actually yielding an array of the things
167+
* that are streamed. For example, if the return type here says
168+
* AsyncGenerator<ChatMessage>, then it's actually AsyncGenerator<ChatMessage[]>.
169+
* This needs to be handled by the caller.
170+
*
171+
* Using unknown for now to make this more explicit
172+
*/
173+
async *streamRequest<T extends keyof FromWebviewProtocol>(
174+
messageType: T,
175+
data: FromWebviewProtocol[T][0],
176+
cancelToken?: AbortSignal,
177+
): AsyncGenerator<
178+
GeneratorYieldType<FromWebviewProtocol[T][1]>[],
179+
GeneratorReturnType<FromWebviewProtocol[T][1]> | undefined
180+
> {
181+
const messageId = uuidv4();
182+
183+
this.post(messageType, data, messageId);
184+
185+
const buffer: GeneratorYieldType<FromWebviewProtocol[T][1]>[] = [];
186+
let index = 0;
187+
let done = false;
188+
let returnVal: GeneratorReturnType<FromWebviewProtocol[T][1]> | undefined =
189+
undefined;
190+
let error: string | null = null;
191+
192+
// This handler receieves individual WebviewMessengerResults
193+
// And pushes them to buffer
194+
const handler = (event: {
195+
data: Message<WebviewProtocolGeneratorMessage<T>>;
196+
}) => {
197+
if (event.data.messageId === messageId) {
198+
const responseData = event.data.data;
199+
if ("error" in responseData) {
200+
error = responseData.error;
201+
return;
202+
// throw new Error(responseData.error);
203+
}
204+
if (responseData.done) {
205+
window.removeEventListener("message", handler);
206+
done = true;
207+
returnVal = responseData.content;
208+
} else {
209+
buffer.push(responseData.content);
210+
}
211+
}
212+
};
213+
window.addEventListener("message", handler);
214+
215+
const handleAbort = () => {
216+
this.post("abort", undefined, messageId);
217+
};
218+
cancelToken?.addEventListener("abort", handleAbort);
219+
220+
try {
221+
while (!done) {
222+
if (error) {
223+
throw error;
224+
}
225+
if (buffer.length > index) {
226+
const chunks = buffer.slice(index);
227+
index = buffer.length;
228+
yield chunks;
229+
}
230+
await new Promise((resolve) => setTimeout(resolve, 50));
231+
}
232+
233+
if (buffer.length > index) {
234+
const chunks = buffer.slice(index);
235+
yield chunks;
236+
}
237+
238+
if (!returnVal) {
239+
return undefined;
240+
}
241+
return returnVal;
242+
} catch (e) {
243+
throw e;
244+
} finally {
245+
cancelToken?.removeEventListener("abort", handleAbort);
246+
}
247+
}
248+
249+
async *llmStreamChat(
250+
msg: ToCoreProtocol["llm/streamChat"][0],
251+
cancelToken: AbortSignal,
252+
): AsyncGenerator<ChatMessage[], PromptLog | undefined> {
253+
const gen = this.streamRequest("llm/streamChat", msg, cancelToken);
254+
255+
let next = await gen.next();
256+
while (!next.done) {
257+
yield next.value;
258+
next = await gen.next();
259+
}
260+
return next.value;
261+
}
262+
}
263+
264+
export const IdeMessengerContext = createContext<IIdeMessenger>(
265+
new IdeMessenger(),
266+
);
267+
268+
export const IdeMessengerProvider: React.FC<{
269+
children: React.ReactNode;
270+
messenger?: IIdeMessenger;
271+
}> = ({ children, messenger = new IdeMessenger() }) => {
272+
return (
273+
<IdeMessengerContext.Provider value={messenger}>
274+
{children}
275+
</IdeMessengerContext.Provider>
276+
);
277+
};

‎gui/src/pages/gui/Chat.test.tsx‎

Lines changed: 54 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,62 @@
11
import { renderWithProviders } from "../../util/test/render";
2-
import { screen } from "@testing-library/dom";
2+
import { screen, waitFor } from "@testing-library/dom";
3+
import { act, fireEvent } from "@testing-library/react";
34
import { Chat } from "./Chat";
45

56
describe("Chat page test", () => {
6-
it("should render", async () => {
7+
it("should render input box", async () => {
78
await renderWithProviders(<Chat />);
8-
99
expect(await screen.findByTestId("continue-input-box")).toBeInTheDocument();
1010
});
11+
12+
it("should be able to toggle modes", async () => {
13+
await renderWithProviders(<Chat />);
14+
expect(screen.getByText("Chat")).toBeInTheDocument();
15+
16+
// Simulate cmd+. keyboard shortcut to toggle modes
17+
act(() => {
18+
document.dispatchEvent(
19+
new KeyboardEvent("keydown", {
20+
key: ".",
21+
metaKey: true, // cmd key on Mac
22+
}),
23+
);
24+
});
25+
26+
// Check that it switched to Edit mode
27+
expect(await screen.findByText("Edit")).toBeInTheDocument();
28+
29+
act(() => {
30+
document.dispatchEvent(
31+
new KeyboardEvent("keydown", {
32+
key: ".",
33+
metaKey: true, // cmd key on Mac
34+
}),
35+
);
36+
});
37+
38+
// Check that it switched to Agent mode
39+
expect(await screen.findByText("Agent")).toBeInTheDocument();
40+
});
41+
42+
it.skip("should send a message and receive a response", async () => {
43+
const { user, container } = await renderWithProviders(<Chat />);
44+
const inputBox = await waitFor(() =>
45+
container.querySelector(".ProseMirror")!.querySelector("p"),
46+
);
47+
expect(inputBox).toBeDefined();
48+
49+
const sendButton = await screen.findByTestId("submit-input-button");
50+
51+
await act(async () => {
52+
// Focus input box
53+
inputBox!.focus();
54+
55+
// Type message
56+
await user.type(inputBox!, "Hello, world!");
57+
58+
sendButton.click();
59+
});
60+
expect(await screen.findByText("Hello, world!")).toBeInTheDocument();
61+
});
1162
});

0 commit comments

Comments
 (0)