diff --git a/src/api/generated/@tanstack/react-query.gen.ts b/src/api/generated/@tanstack/react-query.gen.ts index 03b41ea4..3b2b6a9f 100644 --- a/src/api/generated/@tanstack/react-query.gen.ts +++ b/src/api/generated/@tanstack/react-query.gen.ts @@ -16,6 +16,9 @@ import { v1DeleteWorkspace, v1GetWorkspaceAlerts, v1GetWorkspaceMessages, + v1GetWorkspaceSystemPrompt, + v1SetWorkspaceSystemPrompt, + v1DeleteWorkspaceSystemPrompt, } from "../sdk.gen"; import type { V1CreateWorkspaceData, @@ -29,6 +32,13 @@ import type { V1DeleteWorkspaceResponse, V1GetWorkspaceAlertsData, V1GetWorkspaceMessagesData, + V1GetWorkspaceSystemPromptData, + V1SetWorkspaceSystemPromptData, + V1SetWorkspaceSystemPromptError, + V1SetWorkspaceSystemPromptResponse, + V1DeleteWorkspaceSystemPromptData, + V1DeleteWorkspaceSystemPromptError, + V1DeleteWorkspaceSystemPromptResponse, } from "../types.gen"; type QueryKey = [ @@ -343,3 +353,64 @@ export const v1GetWorkspaceMessagesOptions = ( queryKey: v1GetWorkspaceMessagesQueryKey(options), }); }; + +export const v1GetWorkspaceSystemPromptQueryKey = ( + options: OptionsLegacyParser, +) => [createQueryKey("v1GetWorkspaceSystemPrompt", options)]; + +export const v1GetWorkspaceSystemPromptOptions = ( + options: OptionsLegacyParser, +) => { + return queryOptions({ + queryFn: async ({ queryKey, signal }) => { + const { data } = await v1GetWorkspaceSystemPrompt({ + ...options, + ...queryKey[0], + signal, + throwOnError: true, + }); + return data; + }, + queryKey: v1GetWorkspaceSystemPromptQueryKey(options), + }); +}; + +export const v1SetWorkspaceSystemPromptMutation = ( + options?: Partial>, +) => { + const mutationOptions: UseMutationOptions< + V1SetWorkspaceSystemPromptResponse, + V1SetWorkspaceSystemPromptError, + OptionsLegacyParser + > = { + mutationFn: async (localOptions) => { + const { data } = await v1SetWorkspaceSystemPrompt({ + ...options, + ...localOptions, + throwOnError: true, + }); + return data; + }, + }; + return mutationOptions; +}; + +export const v1DeleteWorkspaceSystemPromptMutation = ( + options?: Partial>, +) => { + const mutationOptions: UseMutationOptions< + V1DeleteWorkspaceSystemPromptResponse, + V1DeleteWorkspaceSystemPromptError, + OptionsLegacyParser + > = { + mutationFn: async (localOptions) => { + const { data } = await v1DeleteWorkspaceSystemPrompt({ + ...options, + ...localOptions, + throwOnError: true, + }); + return data; + }, + }; + return mutationOptions; +}; diff --git a/src/api/generated/sdk.gen.ts b/src/api/generated/sdk.gen.ts index 26c83a75..9df3646e 100644 --- a/src/api/generated/sdk.gen.ts +++ b/src/api/generated/sdk.gen.ts @@ -35,6 +35,15 @@ import type { V1GetWorkspaceMessagesData, V1GetWorkspaceMessagesError, V1GetWorkspaceMessagesResponse, + V1GetWorkspaceSystemPromptData, + V1GetWorkspaceSystemPromptError, + V1GetWorkspaceSystemPromptResponse, + V1SetWorkspaceSystemPromptData, + V1SetWorkspaceSystemPromptError, + V1SetWorkspaceSystemPromptResponse, + V1DeleteWorkspaceSystemPromptData, + V1DeleteWorkspaceSystemPromptError, + V1DeleteWorkspaceSystemPromptResponse, } from "./types.gen"; export const client = createClient(createConfig()); @@ -243,3 +252,58 @@ export const v1GetWorkspaceMessages = ( url: "/api/v1/workspaces/{workspace_name}/messages", }); }; + +/** + * Get Workspace System Prompt + * Get the system prompt for a workspace. + */ +export const v1GetWorkspaceSystemPrompt = < + ThrowOnError extends boolean = false, +>( + options: OptionsLegacyParser, +) => { + return (options?.client ?? client).get< + V1GetWorkspaceSystemPromptResponse, + V1GetWorkspaceSystemPromptError, + ThrowOnError + >({ + ...options, + url: "/api/v1/workspaces/{workspace_name}/system-prompt", + }); +}; + +/** + * Set Workspace System Prompt + */ +export const v1SetWorkspaceSystemPrompt = < + ThrowOnError extends boolean = false, +>( + options: OptionsLegacyParser, +) => { + return (options?.client ?? client).put< + V1SetWorkspaceSystemPromptResponse, + V1SetWorkspaceSystemPromptError, + ThrowOnError + >({ + ...options, + url: "/api/v1/workspaces/{workspace_name}/system-prompt", + }); +}; + +/** + * Delete Workspace System Prompt + */ +export const v1DeleteWorkspaceSystemPrompt = < + ThrowOnError extends boolean = false, +>( + options: OptionsLegacyParser, +) => { + return (options?.client ?? client).delete< + V1DeleteWorkspaceSystemPromptResponse, + V1DeleteWorkspaceSystemPromptError, + ThrowOnError + >({ + ...options, + url: "/api/v1/workspaces/{workspace_name}/system-prompt", + }); +}; diff --git a/src/api/generated/types.gen.ts b/src/api/generated/types.gen.ts index fb0e0f7f..e34a9e1f 100644 --- a/src/api/generated/types.gen.ts +++ b/src/api/generated/types.gen.ts @@ -79,6 +79,10 @@ export type QuestionAnswer = { answer: ChatMessage | null; }; +export type SystemPrompt = { + prompt: string; +}; + export type ValidationError = { loc: Array; msg: string; @@ -166,3 +170,34 @@ export type V1GetWorkspaceMessagesData = { export type V1GetWorkspaceMessagesResponse = Array; export type V1GetWorkspaceMessagesError = HTTPValidationError; + +export type V1GetWorkspaceSystemPromptData = { + path: { + workspace_name: string; + }; +}; + +export type V1GetWorkspaceSystemPromptResponse = SystemPrompt; + +export type V1GetWorkspaceSystemPromptError = HTTPValidationError; + +export type V1SetWorkspaceSystemPromptData = { + body: SystemPrompt; + path: { + workspace_name: string; + }; +}; + +export type V1SetWorkspaceSystemPromptResponse = void; + +export type V1SetWorkspaceSystemPromptError = HTTPValidationError; + +export type V1DeleteWorkspaceSystemPromptData = { + path: { + workspace_name: string; + }; +}; + +export type V1DeleteWorkspaceSystemPromptResponse = void; + +export type V1DeleteWorkspaceSystemPromptError = HTTPValidationError; diff --git a/src/api/openapi.json b/src/api/openapi.json index 744fc9d2..9e0dd2ba 100644 --- a/src/api/openapi.json +++ b/src/api/openapi.json @@ -373,6 +373,119 @@ } } } + }, + "/api/v1/workspaces/{workspace_name}/system-prompt": { + "get": { + "tags": ["CodeGate API", "Workspaces"], + "summary": "Get Workspace System Prompt", + "description": "Get the system prompt for a workspace.", + "operationId": "v1_get_workspace_system_prompt", + "parameters": [ + { + "name": "workspace_name", + "in": "path", + "required": true, + "schema": { + "type": "string", + "title": "Workspace Name" + } + } + ], + "responses": { + "200": { + "description": "Successful Response", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/SystemPrompt" + } + } + } + }, + "422": { + "description": "Validation Error", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + } + } + } + }, + "put": { + "tags": ["CodeGate API", "Workspaces"], + "summary": "Set Workspace System Prompt", + "operationId": "v1_set_workspace_system_prompt", + "parameters": [ + { + "name": "workspace_name", + "in": "path", + "required": true, + "schema": { + "type": "string", + "title": "Workspace Name" + } + } + ], + "requestBody": { + "required": true, + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/SystemPrompt" + } + } + } + }, + "responses": { + "204": { + "description": "Successful Response" + }, + "422": { + "description": "Validation Error", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + } + } + } + }, + "delete": { + "tags": ["CodeGate API", "Workspaces"], + "summary": "Delete Workspace System Prompt", + "operationId": "v1_delete_workspace_system_prompt", + "parameters": [ + { + "name": "workspace_name", + "in": "path", + "required": true, + "schema": { + "type": "string", + "title": "Workspace Name" + } + } + ], + "responses": { + "204": { + "description": "Successful Response" + }, + "422": { + "description": "Validation Error", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + } + } + } + } } }, "components": { @@ -652,6 +765,17 @@ "title": "QuestionAnswer", "description": "Represents a question and answer pair." }, + "SystemPrompt": { + "properties": { + "prompt": { + "type": "string", + "title": "Prompt" + } + }, + "type": "object", + "required": ["prompt"], + "title": "SystemPrompt" + }, "ValidationError": { "properties": { "loc": { diff --git a/src/features/workspace-system-prompt/components/__tests__/system-prompt-editor.test.tsx b/src/features/workspace-system-prompt/components/__tests__/system-prompt-editor.test.tsx index 65bf7d42..c5443804 100644 --- a/src/features/workspace-system-prompt/components/__tests__/system-prompt-editor.test.tsx +++ b/src/features/workspace-system-prompt/components/__tests__/system-prompt-editor.test.tsx @@ -2,7 +2,8 @@ import { render, waitFor } from "@/lib/test-utils"; import { expect, test } from "vitest"; import { SystemPromptEditor } from "../system-prompt-editor"; import userEvent from "@testing-library/user-event"; -import * as POST_MODULE from "../../lib/post-system-prompt"; +import { server } from "@/mocks/msw/node"; +import { http, HttpResponse } from "msw"; vi.mock("../../lib/post-system-prompt"); @@ -19,25 +20,38 @@ vi.mock("@monaco-editor/react", () => { return { default: FakeEditor }; }); -const renderComponent = () => render(); +const renderComponent = () => + render(); test("can update system prompt", async () => { - // eslint-disable-next-line @typescript-eslint/no-unused-vars - const POST_PROMPT_MOCK = vi.fn(async (_) => Promise.resolve()); - - vi.mocked(POST_MODULE).postSystemPrompt = POST_PROMPT_MOCK; + server.use( + http.get("*/api/v1/workspaces/:name/system-prompt", () => { + return HttpResponse.json({ prompt: "initial prompt from server" }); + }), + ); const { getByRole } = renderComponent(); + await waitFor(() => { + expect(getByRole("textbox")).toBeVisible(); + }); + const input = getByRole("textbox"); - expect(input).toBeVisible(); + expect(input).toHaveTextContent("initial prompt from server"); await userEvent.clear(input); - await userEvent.type(input, "foo bar 123"); - expect(input).toHaveTextContent("foo bar 123"); + await userEvent.type(input, "new prompt from test"); + expect(input).toHaveTextContent("new prompt from test"); await userEvent.click(getByRole("button", { name: /save changes/i })); + + server.use( + http.get("*/api/v1/workspaces/:name/system-prompt", () => { + return HttpResponse.json({ prompt: "new prompt from test" }); + }), + ); + await waitFor(() => { - expect(POST_PROMPT_MOCK).toBeCalledWith("foo bar 123"); + expect(input).toHaveTextContent("new prompt from test"); }); }); diff --git a/src/features/workspace-system-prompt/components/system-prompt-editor.tsx b/src/features/workspace-system-prompt/components/system-prompt-editor.tsx index 745c1b53..a92cb3b8 100644 --- a/src/features/workspace-system-prompt/components/system-prompt-editor.tsx +++ b/src/features/workspace-system-prompt/components/system-prompt-editor.tsx @@ -6,26 +6,39 @@ import { CardFooter, DarkModeContext, LinkButton, + Loader, Text, } from "@stacklok/ui-kit"; import { Dispatch, SetStateAction, + useCallback, useContext, useEffect, + useMemo, useState, } from "react"; -import { usePostSystemPrompt } from "../hooks/use-post-system-prompt"; +import { usePostSystemPrompt } from "../hooks/use-set-system-prompt"; import { Check } from "lucide-react"; import { twMerge } from "tailwind-merge"; +import { + V1GetWorkspaceSystemPromptData, + V1GetWorkspaceSystemPromptResponse, + V1SetWorkspaceSystemPromptData, +} from "@/api/generated"; +import { useGetSystemPrompt } from "../hooks/use-get-system-prompt"; +import { + QueryCacheNotifyEvent, + QueryClient, + useQueryClient, +} from "@tanstack/react-query"; +import { v1GetWorkspaceSystemPromptQueryKey } from "@/api/generated/@tanstack/react-query.gen"; type DarkModeContextValue = { preference: "dark" | "light" | null; override: "dark" | "light" | null; }; -const DEFAULT_VALUE = `You are a security expert conducting a thorough code review.\nIdentify potential security vulnerabilities, suggest improvements, and explain security best practices.`; - function inferDarkMode( contextValue: | null @@ -53,16 +66,120 @@ function useSavedStatus() { return { saved, setSaved }; } -export function SystemPromptEditor({ className }: { className?: string }) { +function EditorLoadingUI() { + return ( + // arbitrary value to match the monaco editor height + // eslint-disable-next-line tailwindcss/no-unnecessary-arbitrary-value +
+ +
+ ); +} + +function isGetSystemPromptQuery( + queryKey: unknown, + options: V1GetWorkspaceSystemPromptData, +): boolean { + return ( + Array.isArray(queryKey) && + queryKey[0]._id === v1GetWorkspaceSystemPromptQueryKey(options)[0]?._id + ); +} + +function getPromptFromNotifyEvent(event: QueryCacheNotifyEvent): string | null { + if ("action" in event === false || "data" in event.action === false) + return null; + return ( + (event.action.data as V1GetWorkspaceSystemPromptResponse | undefined | null) + ?.prompt ?? null + ); +} + +function usePromptValue({ + initialValue, + options, + queryClient, +}: { + initialValue: string; + options: V1GetWorkspaceSystemPromptData; + queryClient: QueryClient; +}) { + const [value, setValue] = useState(initialValue); + + // Subscribe to changes in the workspace system prompt value in the query cache + useEffect(() => { + const queryCache = queryClient.getQueryCache(); + const unsubscribe = queryCache.subscribe((event) => { + if ( + event.type === "updated" && + event.action.type === "success" && + isGetSystemPromptQuery(event.query.options.queryKey, options) + ) { + const prompt: string | null = getPromptFromNotifyEvent(event); + if (prompt === value || prompt === null) return; + + setValue(prompt); + } + }); + + return () => { + return unsubscribe(); + }; + }, [options, queryClient, value]); + + return { value, setValue }; +} + +export function SystemPromptEditor({ + className, + workspaceName, +}: { + className?: string; + workspaceName: string; +}) { const context = useContext(DarkModeContext); const theme: Theme = inferDarkMode(context); - const [value, setValue] = useState(() => DEFAULT_VALUE); + const options: V1GetWorkspaceSystemPromptData & + Omit = useMemo( + () => ({ + path: { workspace_name: workspaceName }, + }), + [workspaceName], + ); + + const queryClient = useQueryClient(); - const { mutate, isPending } = usePostSystemPrompt(); + const { data: systemPromptResponse, isPending: isGetPromptPending } = + useGetSystemPrompt(options); + const { mutate, isPending: isMutationPending } = usePostSystemPrompt(options); + + const { setValue, value } = usePromptValue({ + initialValue: systemPromptResponse?.prompt ?? "", + options, + queryClient, + }); const { saved, setSaved } = useSavedStatus(); + const handleSubmit = useCallback( + (value: string) => { + mutate( + { ...options, body: { prompt: value } }, + { + onSuccess: () => { + queryClient.invalidateQueries({ + queryKey: v1GetWorkspaceSystemPromptQueryKey(options), + refetchType: "all", + }); + setSaved(true); + }, + }, + ); + }, + [mutate, options, queryClient, setSaved], + ); + return ( @@ -72,18 +189,21 @@ export function SystemPromptEditor({ className }: { className?: string }) { save time & tokens.
- setValue(v ?? "")} - height="20rem" - defaultLanguage="Markdown" - theme={theme} - className="bg-base" - defaultValue="" - /> + {isGetPromptPending ? ( + + ) : ( + setValue(v ?? "")} + height="20rem" + defaultLanguage="Markdown" + theme={theme} + className="bg-base" + /> + )}
@@ -91,13 +211,9 @@ export function SystemPromptEditor({ className }: { className?: string }) { Cancel