2 files changed, 22 insertions, 16 deletions
diff --git a/extension/react-app/src/components/CodeBlock.tsx b/extension/react-app/src/components/CodeBlock.tsx
index 9438ab23..98760fbd 100644
--- a/extension/react-app/src/components/CodeBlock.tsx
+++ b/extension/react-app/src/components/CodeBlock.tsx
@@ -13,20 +13,19 @@ const StyledPre = styled.pre`
   border: 0.5px solid gray;
   border-radius: ${defaultBorderRadius};
   background-color: ${secondaryDark};
-  padding: 8px;
-  padding-top: 14px;
-  padding-bottom: 16px;
+  position: relative;
+`;
+
+const CopyButtonDiv = styled.div`
+  position: absolute;
+  top: 4px;
+  right: 4px;
 `;
 
 const StyledCopyButton = styled.button<{ visible: boolean }>`
-  /* position: relative; */
-  float: right;
   border: none;
-  background-color: ${secondaryDark};
+  background-color: transparent;
   cursor: pointer;
-  padding: 0;
-  /* margin: 4px; */
-  margin-top: -6px;
 
   visibility: ${(props) => (props.visible ? "visible" : "hidden")};
 `;
@@ -77,11 +76,13 @@ function CodeBlock(props: { language?: string; children: string }) {
         setHovered(false);
       }}
     >
-      <CopyButton
-        visible={hovered}
-        textToCopy={(props.children as any).props.children[0]}
-      />
-      <StyledCode>{props.children}</StyledCode>
+      <CopyButtonDiv>
+        <CopyButton
+          visible={hovered}
+          textToCopy={(props.children as any).props.children[0]}
+        />
+      </CopyButtonDiv>
+      <StyledCode language={props.language}>{props.children}</StyledCode>
     </StyledPre>
   );
 }
diff --git a/extension/react-app/src/components/StyledCode.tsx b/extension/react-app/src/components/StyledCode.tsx
index 66135d65..485c70b6 100644
--- a/extension/react-app/src/components/StyledCode.tsx
+++ b/extension/react-app/src/components/StyledCode.tsx
@@ -1,12 +1,17 @@
 import { Prism as SyntaxHighlighter } from "react-syntax-highlighter";
-import { synthwave84 } from "react-syntax-highlighter/dist/esm/styles/prism";
+import { vscDarkPlus as highlightStyle } from "react-syntax-highlighter/dist/esm/styles/prism";
 
 interface StyledCodeProps {
   children: string;
+  language?: string;
 }
 
 const StyledCode = (props: StyledCodeProps) => (
-  <SyntaxHighlighter style={synthwave84}>
+  <SyntaxHighlighter
+    customStyle={{ margin: "0" }}
+    style={highlightStyle}
+    language={props.language || "python"}
+  >
     {(props.children as any).props.children}
   </SyntaxHighlighter>
 );