Skip to content

Commit

Permalink
fix(studio): codemod learn fixes with ws (#1195)
Browse files Browse the repository at this point in the history
  • Loading branch information
@grzim
grzim authored Aug 7, 2024
1 parent 4dfa432 commit c3b48e2
Show file tree
Hide file tree
Showing 2 changed files with 152 additions and 121 deletions.
54 changes: 30 additions & 24 deletions apps/frontend/app/(website)/studio/features/modGPT/useAiService/codemodAI/useCodemodAI.tsx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -75,28 +75,29 @@ export const useCodemodAI = ({
}
};

useEffect(() => {
const connectWebSocket = async () => {
const token = await getToken();
if (token) {
const websocket = new WebSocket(env.NEXT_PUBLIC_WS_URL as string);
setWs(websocket);
setIsWsConnected(true);
const connectWebSocket = async () => {
const token = await getToken();
if (token) {
const websocket = new WebSocket(env.NEXT_PUBLIC_WS_URL as string);
setWs(websocket);

websocket.onopen = () =>
console.info("WebSocket connection established");
websocket.onmessage = (event) => onMessage(JSON.parse(event.data));
websocket.onerror = handleError;
websocket.onclose = () => setIsWsConnected(false);
websocket.onopen = () => {
console.info("WebSocket connection established");
setIsWsConnected(true);
};
websocket.onmessage = (event) => onMessage(JSON.parse(event.data));
websocket.onerror = handleError;
websocket.onclose = () => setIsWsConnected(false);

return () => {
websocket.close();
setIsWsConnected(false);
setServiceBusy(false);
};
}
};
return () => {
websocket.close();
setIsWsConnected(false);
setServiceBusy(false);
};
}
};

useEffect(() => {
connectWebSocket();
}, []);

Expand All @@ -120,10 +121,12 @@ export const useCodemodAI = ({
}
};

const startIterativeCodemodGeneration = async () => {
const startIterativeCodemodGeneration = async (
content = `Generate codemod with AI`,
) => {
if (isEnvPrepared) {
setWsMessage({
content: `Generate codemod with AI`,
content,
role: "user",
id: Date.now().toString(),
});
Expand All @@ -138,15 +141,18 @@ export const useCodemodAI = ({
};

useEffect(() => {
if (command === LEARN_KEY) {
startIterativeCodemodGeneration();
if (command === LEARN_KEY && isWsConnected) {
startIterativeCodemodGeneration(
"Codemod learn: generate codemod with AI",
);
setCurrentCommand(null);
}
}, [command, isEnvPrepared]);
}, [command, isEnvPrepared, isWsConnected]);

return {
stopCodemodAi: () => {
ws?.close();
connectWebSocket();
setIsWsConnected(false);
setServiceBusy(false);
},
Expand Down
219 changes: 122 additions & 97 deletions apps/frontend/app/(website)/studio/features/share/useSharedState.ts
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
"use client";
import { LEARN_KEY } from "@/constants";
import type { KnownEngines } from "@codemod-com/utilities";
import { getCodeDiff } from "@studio/api/getCodeDiff";
Expand Down Expand Up @@ -27,50 +26,8 @@ const decodeNullable = (value: string | null): string | null => {
return value;
}
};
const getState = async () => {
const searchParams = new URLSearchParams(window.location.search);

const csc = searchParams.get(SEARCH_PARAMS_KEYS.COMPRESSED_SHAREABLE_CODEMOD);

const engine =
decodeNullable(searchParams.get(SEARCH_PARAMS_KEYS.ENGINE)) ??
("jscodeshift" as KnownEngines);
const diffId = searchParams.get(SEARCH_PARAMS_KEYS.DIFF_ID);
const codemodSource = decodeNullable(
searchParams.get(SEARCH_PARAMS_KEYS.CODEMOD_SOURCE),
);
const codemodName = decodeNullable(
searchParams.get(SEARCH_PARAMS_KEYS.CODEMOD_NAME),
);

const command = searchParams.get(SEARCH_PARAMS_KEYS.COMMAND);
const iv = searchParams.get(SEARCH_PARAMS_KEYS.IV);
const isLearn = command === LEARN_KEY;

if (isLearn) {
if (!engine || !diffId || !iv) {
return;
}

const snippets = await getCodeDiff({ diffId, iv });

if (!snippets) {
return;
}

const editors = [toInitialStates({ ...snippets, name: "Learn" })];
searchParams.delete(SEARCH_PARAMS_KEYS.COMMAND);
return {
engine,
editors,
codemodSource: "",
codemodName,
command: LEARN_KEY,
};
}

if (csc === null) return;

const parseShareableCodemodFromParams = (csc: string) => {
try {
const encryptedString = window.atob(
csc.replaceAll("-", "+").replaceAll("_", "/"),
Expand All @@ -83,49 +40,77 @@ const getState = async () => {
const uint8Array = Uint8Array.from(numberArray);

const decryptedString = inflate(uint8Array, { to: "string" });
const shareableCodemod = parseShareableCodemod(JSON.parse(decryptedString));

const getMultipleEditors = ({
before,
after,
names,
}: {
before: string[];
after: string[];
names: string[];
}) => {
const zipit = zipWith((before, after) => ({ before, after }));
const zipitMore = zipWith(({ before, after }, name) => ({
before,
after,
name: name ?? "test",
}));
return zipitMore(zipit(before, after), names);
};

const editors = shareableCodemod.bm
? getMultipleEditors({
before: shareableCodemod.bm.split("__codemod_splitter__"),
after: shareableCodemod.am.split("__codemod_splitter__"),
names: shareableCodemod.nm.split("__codemod_splitter__"),
})
: [
{
name: "test 1",
before: shareableCodemod.b ?? "",
after: shareableCodemod.a ?? "",
},
];
return {
engine: shareableCodemod.e ?? "jscodeshift",
editors: editors.map(toInitialStates),
codemodSource: shareableCodemod.c ?? "",
codemodName: shareableCodemod.n ?? null,
command: null,
};
return parseShareableCodemod(JSON.parse(decryptedString));
} catch (error) {
console.error(error);
return null;
}
};

const getMultipleEditors = ({
before,
after,
names,
}: {
before: string[];
after: string[];
names: string[];
}) => {
const zipit = zipWith((before, after) => ({ before, after }));
const zipitMore = zipWith(({ before, after }, name) => ({
before,
after,
name: name ?? "test",
}));
return zipitMore(zipit(before, after), names);
};

const createStateFromCodemod = (shareableCodemod: any) => {
const editors = shareableCodemod.bm
? getMultipleEditors({
before: shareableCodemod.bm.split("__codemod_splitter__"),
after: shareableCodemod.am.split("__codemod_splitter__"),
names: shareableCodemod.nm.split("__codemod_splitter__"),
})
: [
{
name: "test 1",
before: shareableCodemod.b ?? "",
after: shareableCodemod.a ?? "",
},
];
return {
engine: shareableCodemod.e ?? "jscodeshift",
editors: editors.map(toInitialStates),
codemodSource: shareableCodemod.c ?? "",
codemodName: shareableCodemod.n ?? null,
command: null,
};
};

const getState = () => {
const searchParams = new URLSearchParams(window.location.search);

const csc = searchParams.get(SEARCH_PARAMS_KEYS.COMPRESSED_SHAREABLE_CODEMOD);
if (csc !== null) {
const shareableCodemod = parseShareableCodemodFromParams(csc);
if (shareableCodemod) {
return createStateFromCodemod(shareableCodemod);
}
}

const engine = decodeNullable(
searchParams.get(SEARCH_PARAMS_KEYS.ENGINE),
) as KnownEngines;
const diffId = searchParams.get(SEARCH_PARAMS_KEYS.DIFF_ID);
const codemodSource = decodeNullable(
searchParams.get(SEARCH_PARAMS_KEYS.CODEMOD_SOURCE),
);
const codemodName = decodeNullable(
searchParams.get(SEARCH_PARAMS_KEYS.CODEMOD_NAME),
);

const command = searchParams.get(SEARCH_PARAMS_KEYS.COMMAND);

const someSearchParamsSet = [
engine,
Expand All @@ -135,32 +120,72 @@ const getState = async () => {
command,
].some((s) => s !== null);

if (someSearchParamsSet)
if (someSearchParamsSet) {
return {
engine: engine ?? "jscodeshift",
editors: [getEmptyTestCase()],
codemodSource: codemodSource ?? "",
codemodName: codemodName ?? "",
command:
command === LEARN_KEY || command === "accessTokenRequested"
command === "learn" || command === "accessTokenRequested"
? command
: null,
};
}
};

const handleCodemodLearnState = async (
searchParams: URLSearchParams,
snippetStore: any,
setEngine: any,
setCurrentCommand: any,
) => {
try {
const engine = (searchParams.get(SEARCH_PARAMS_KEYS.ENGINE) ??
"jscodeshift") as KnownEngines;
const diffId = searchParams.get(SEARCH_PARAMS_KEYS.DIFF_ID);
const iv = searchParams.get(SEARCH_PARAMS_KEYS.IV);
if (!engine || !diffId || !iv) {
return;
}

const snippets = await getCodeDiff({ diffId, iv });

if (!snippets) {
return;
}

snippetStore.setBeforeSnippet(snippets.before);
snippetStore.setAfterSnippet(snippets.after);
searchParams.delete(SEARCH_PARAMS_KEYS.COMMAND);
setEngine(engine);
setCurrentCommand(LEARN_KEY);
} catch (err) {
console.error(err);
}
};

const useCodemodLearn = () => {};
export const useSharedState = () => {
const { setInitialState } = useSnippetsStore();
const { setInitialState, getSelectedEditors, setEngine } = useSnippetsStore();
const { setContent, setCurrentCommand } = useModStore();
const searchParams = new URLSearchParams(window.location.search);
const command = searchParams.get(SEARCH_PARAMS_KEYS.COMMAND);

useEffect(() => {
const setState = async () => {
const initialState = await getState();
if (initialState) {
setInitialState(initialState);
setContent(initialState.codemodSource);
setCurrentCommand(initialState?.command);
}
};
setState();
const initialState = getState();
const snippetStore = getSelectedEditors();
if (command === LEARN_KEY) {
handleCodemodLearnState(
searchParams,
snippetStore,
setEngine,
setCurrentCommand,
);
return;
}
if (initialState) {
setInitialState(initialState);
setContent(initialState.codemodSource);
}
}, []);
};

0 comments on commit c3b48e2

Please sign in to comment.