3 files changed, 63 insertions, 23 deletions
diff --git a/extension/react-app/src/components/CodeBlock.tsx b/extension/react-app/src/components/CodeBlock.tsx
index e0336554..eedae3fb 100644
--- a/extension/react-app/src/components/CodeBlock.tsx
+++ b/extension/react-app/src/components/CodeBlock.tsx
@@ -11,6 +11,7 @@ const StyledPre = styled.pre`
   border: 1px solid gray;
   border-radius: ${defaultBorderRadius};
   background-color: ${vscBackground};
+  padding: 8px;
 `;
 
 const StyledCode = styled.code`
diff --git a/extension/react-app/src/components/StepContainer.tsx b/extension/react-app/src/components/StepContainer.tsx
index 3408053b..ca142b06 100644
--- a/extension/react-app/src/components/StepContainer.tsx
+++ b/extension/react-app/src/components/StepContainer.tsx
@@ -7,6 +7,7 @@ import {
   vscBackground,
   GradientBorder,
   vscBackgroundTransparent,
+  HeaderButton,
 } from ".";
 import { RangeInFile, FileEdit } from "../../../src/client";
 import CodeBlock from "./CodeBlock";
@@ -15,7 +16,7 @@ import SubContainer from "./SubContainer";
 import {
   ChevronDown,
   ChevronRight,
-  Backward,
+  XMark,
   ArrowPath,
 } from "@styled-icons/heroicons-outline";
 import { HistoryNode } from "../../../schema/HistoryNode";
@@ -31,6 +32,7 @@ interface StepContainerProps {
   onRefinement: (input: string) => void;
   onUserInput: (input: string) => void;
   onRetry: () => void;
+  onDelete: () => void;
   open?: boolean;
 }
 
@@ -54,8 +56,10 @@ const HeaderDiv = styled.div<{ error: boolean }>`
   background-color: ${(props) =>
     props.error ? "#522" : vscBackgroundTransparent};
   display: grid;
-  grid-template-columns: 1fr auto;
+  grid-template-columns: 1fr auto auto;
+  grid-gap: 8px;
   align-items: center;
+  padding-right: 8px;
 `;
 
 const ContentDiv = styled.div`
@@ -64,26 +68,21 @@ const ContentDiv = styled.div`
   background-color: ${vscBackground};
 `;
 
-const HeaderButton = styled.button`
-  background-color: transparent;
-  border: 1px solid white;
-  border-radius: ${defaultBorderRadius};
-  padding: 2px;
-  cursor: pointer;
-  color: white;
-
-  &:hover {
-    background-color: white;
-    color: black;
-  }
-`;
-
 const OnHoverDiv = styled.div`
   text-align: center;
   padding: 10px;
   animation: ${appear} 0.3s ease-in-out;
 `;
 
+const MarkdownPre = styled.pre`
+  background-color: ${secondaryDark};
+  padding: 10px;
+  border-radius: ${defaultBorderRadius};
+  border: 0.5px solid white;
+`;
+
+const MarkdownCode = styled.code``;
+
 function StepContainer(props: StepContainerProps) {
   const [open, setOpen] = useState(
     typeof props.open === "undefined" ? true : props.open
@@ -152,18 +151,28 @@ function StepContainer(props: StepContainerProps) {
               <Backward size="1.6em" onClick={props.onReverse}></Backward>
             </HeaderButton> */}
 
-            {props.historyNode.observation?.error ? (
+            <>
               <HeaderButton
                 onClick={(e) => {
                   e.stopPropagation();
-                  props.onRetry();
+                  props.onDelete();
                 }}
               >
-                <ArrowPath size="1.6em" onClick={props.onRetry}></ArrowPath>
+                <XMark size="1.6em" onClick={props.onDelete} />
               </HeaderButton>
-            ) : (
-              <></>
-            )}
+              {props.historyNode.observation?.error ? (
+                <HeaderButton
+                  onClick={(e) => {
+                    e.stopPropagation();
+                    props.onRetry();
+                  }}
+                >
+                  <ArrowPath size="1.6em" onClick={props.onRetry} />
+                </HeaderButton>
+              ) : (
+                <></>
+              )}
+            </>
           </HeaderDiv>
         </GradientBorder>
         <ContentDiv hidden={!open}>
@@ -182,7 +191,19 @@ function StepContainer(props: StepContainerProps) {
               {props.historyNode.observation.error as string}
             </pre>
           ) : (
-            <ReactMarkdown key={1} className="overflow-scroll">
+            <ReactMarkdown
+              key={1}
+              className="overflow-scroll"
+              components={{
+                pre: ({ node, ...props }) => {
+                  return (
+                    <CodeBlock
+                      children={props.children[0] as string}
+                    ></CodeBlock>
+                  );
+                },
+              }}
+            >
               {props.historyNode.step.description as any}
             </ReactMarkdown>
           )}
diff --git a/extension/react-app/src/components/index.ts b/extension/react-app/src/components/index.ts
index 4966f3e8..525989af 100644
--- a/extension/react-app/src/components/index.ts
+++ b/extension/react-app/src/components/index.ts
@@ -143,3 +143,21 @@ export const appear = keyframes`
         transform: translateY(0px);
     }
 `;
+
+export const HeaderButton = styled.button`
+  background-color: transparent;
+  border: 1px solid white;
+  border-radius: ${defaultBorderRadius};
+  cursor: pointer;
+  color: white;
+
+  &:hover {
+    background-color: white;
+    color: black;
+  }
+  display: flex;
+  align-items: center;
+  justify-content: center;
+  gap: 4px;
+  padding: 1px;
+`;