diff --git a/app/components/pipeline/GeneratePipeline.tsx b/app/components/pipeline/GeneratePipeline.tsx
index 91f3821..3ffedb7 100644
--- a/app/components/pipeline/GeneratePipeline.tsx
+++ b/app/components/pipeline/GeneratePipeline.tsx
@@ -24,7 +24,7 @@ const tabClass = clsx(
 );
 
 export default function GeneratePipeline() {
-  const litlytics = useLitlytics();
+  const { litlytics, pipeline, setPipeline, pipelineStatus } = useLitlytics();
   const [selectedTab, setSelectedTab] = useState(0);
   const [isOpen, setIsOpen] = useState(false);
   const [loading, setLoading] = useState(false);
@@ -36,8 +36,8 @@ export default function GeneratePipeline() {
       setError(undefined);
 
       // generate plan from LLM
-      await litlytics.generatePipeline();
-
+      const newPipeline = await litlytics.generatePipeline();
+      setPipeline(newPipeline);
       setLoading(false);
       setSelectedTab(1);
     } catch (err) {
@@ -49,9 +49,9 @@ export default function GeneratePipeline() {
   const closeDialog = () => {
     if (
       loading ||
-      litlytics.pipelineStatus.status === 'refine' ||
-      litlytics.pipelineStatus.status === 'step' ||
-      litlytics.pipelineStatus.status === 'sourcing'
+      pipelineStatus.status === 'refine' ||
+      pipelineStatus.status === 'step' ||
+      pipelineStatus.status === 'sourcing'
     ) {
       return;
     }
@@ -65,8 +65,8 @@ export default function GeneratePipeline() {
         outline
         onClick={() => setIsOpen(true)}
         disabled={
-          litlytics.pipelineStatus.status === 'sourcing' ||
-          litlytics.pipelineStatus.status === 'step'
+          pipelineStatus.status === 'sourcing' ||
+          pipelineStatus.status === 'step'
         }
         className="border-sky-500 dark:border-sky-500"
       >
@@ -83,7 +83,7 @@ export default function GeneratePipeline() {
           <TabGroup selectedIndex={selectedTab} onChange={setSelectedTab}>
             <TabList className="flex gap-4">
               <Tab className={tabClass}>Plan</Tab>
-              {Boolean(litlytics.pipeline.pipelinePlan?.length) && (
+              {Boolean(pipeline.pipelinePlan?.length) && (
                 <Tab className={tabClass}>Refine plan</Tab>
               )}
             </TabList>
@@ -96,9 +96,10 @@ export default function GeneratePipeline() {
                       rows={5}
                       name="task"
                       placeholder="Task description"
-                      value={litlytics.pipeline.pipelineDescription}
+                      value={pipeline.pipelineDescription ?? ''}
                       onChange={(e) => {
-                        litlytics.setPipeline({
+                        setPipeline({
+                          ...pipeline,
                           pipelineDescription: e.target.value,
                         });
                       }}
diff --git a/app/components/pipeline/PipelineBuilder.tsx b/app/components/pipeline/PipelineBuilder.tsx
index 5632499..a3a4a7c 100644
--- a/app/components/pipeline/PipelineBuilder.tsx
+++ b/app/components/pipeline/PipelineBuilder.tsx
@@ -5,12 +5,12 @@ import { SourceNode } from './nodes/source/SourceNode';
 import { StepNode } from './nodes/StepNode';
 
 export function PipelineBuilder() {
-  const litlytics = useLitlytics();
+  const { pipeline } = useLitlytics();
 
   return (
     <Background>
       <SourceNode />
-      {litlytics.pipeline.steps
+      {pipeline.steps
         .sort((a, b) => (a.connectsTo.includes(b.id) ? -1 : 1))
         .map((step) => (
           <StepNode data={step} key={step.id} />
diff --git a/app/components/pipeline/RefinePipeline.tsx b/app/components/pipeline/RefinePipeline.tsx
index 3748ebf..1280dce 100644
--- a/app/components/pipeline/RefinePipeline.tsx
+++ b/app/components/pipeline/RefinePipeline.tsx
@@ -6,7 +6,13 @@ import { Spinner } from '~/components/Spinner';
 import { useLitlytics } from '~/store/WithLitLytics';
 
 export function RefinePipeline({ hide }: { hide: () => void }) {
-  const litlytics = useLitlytics();
+  const {
+    litlytics,
+    pipeline,
+    setPipeline,
+    pipelineStatus,
+    setPipelineStatus,
+  } = useLitlytics();
   const [error, setError] = useState<Error>();
   const [refine, setRefine] = useState(``);
   const [progress, setProgress] = useState('');
@@ -18,40 +24,46 @@ export function RefinePipeline({ hide }: { hide: () => void }) {
 
     try {
       // generate plan from LLM
-      litlytics.setPipelineStatus({ status: 'refine' });
-      await litlytics.refinePipeline({
+      setPipelineStatus({ status: 'refine' });
+      const newPipeline = await litlytics.refinePipeline({
         refineRequest: refine,
       });
-      litlytics.setPipelineStatus({ status: 'init' });
+      setPipeline(newPipeline);
+      setPipelineStatus({ status: 'init' });
       setRefine('');
     } catch (err) {
       setError(err as Error);
-      litlytics.setPipelineStatus({ status: 'error' });
+      setPipelineStatus({ status: 'error' });
     }
   };
 
   const doCreate = async () => {
     try {
+      setPipelineStatus({ status: 'sourcing' });
       // generate plan from LLM
-      await litlytics.pipelineFromText(({ step, totalSteps }) => {
-        if (step > totalSteps) {
-          setProgress('');
-          return;
-        }
+      const newPipeline = await litlytics.pipelineFromText(
+        ({ step, totalSteps }) => {
+          if (step > totalSteps) {
+            setProgress('');
+            return;
+          }
 
-        setProgress(`Generating steps: ${step} / ${totalSteps}`);
-      });
+          setProgress(`Generating steps: ${step} / ${totalSteps}`);
+        }
+      );
+      setPipeline(newPipeline);
+      setPipelineStatus({ status: 'done' });
       hide();
     } catch (err) {
       setError(err as Error);
-      litlytics.setPipelineStatus({ status: 'error' });
+      setPipelineStatus({ status: 'error' });
     }
   };
 
   const inProgress =
-    litlytics.pipelineStatus.status === 'refine' ||
-    litlytics.pipelineStatus.status === 'sourcing' ||
-    litlytics.pipelineStatus.status === 'step';
+    pipelineStatus.status === 'refine' ||
+    pipelineStatus.status === 'sourcing' ||
+    pipelineStatus.status === 'step';
 
   return (
     <div className="w-full h-full p-4 prose prose-sm prose-no-nr dark:prose-invert ">
@@ -59,7 +71,7 @@ export function RefinePipeline({ hide }: { hide: () => void }) {
         <h1 className="m-0">Suggested pipeline:</h1>
       </div>
 
-      <CustomMarkdown>{`${litlytics.pipeline.pipelinePlan}`}</CustomMarkdown>
+      <CustomMarkdown>{`${pipeline.pipelinePlan}`}</CustomMarkdown>
       <div className="flex gap-1">
         <Textarea
           rows={2}
@@ -69,7 +81,7 @@ export function RefinePipeline({ hide }: { hide: () => void }) {
           onChange={(e) => setRefine(e.target.value)}
         />
         <Button onClick={doRefine} disabled={inProgress}>
-          {litlytics.pipelineStatus.status === 'refine' && (
+          {pipelineStatus.status === 'refine' && (
             <Spinner className="h-5 w-5" />
           )}
           Refine
@@ -84,14 +96,14 @@ export function RefinePipeline({ hide }: { hide: () => void }) {
 
       <div className="flex mt-8">
         <Button onClick={doCreate} disabled={inProgress}>
-          {(litlytics.pipelineStatus.status === 'sourcing' ||
-            litlytics.pipelineStatus.status === 'step') && (
+          {(pipelineStatus.status === 'sourcing' ||
+            pipelineStatus.status === 'step') && (
             <div className="flex items-center">
               <Spinner className="h-5 w-5" />
             </div>
           )}
-          {litlytics.pipelineStatus.status === 'sourcing' ||
-          litlytics.pipelineStatus.status === 'step'
+          {pipelineStatus.status === 'sourcing' ||
+          pipelineStatus.status === 'step'
             ? progress.length > 0
               ? progress
               : `Generating pipeline...`
diff --git a/app/components/pipeline/nodes/NodeConnector.tsx b/app/components/pipeline/nodes/NodeConnector.tsx
index b56a67b..d59802d 100644
--- a/app/components/pipeline/nodes/NodeConnector.tsx
+++ b/app/components/pipeline/nodes/NodeConnector.tsx
@@ -1,7 +1,6 @@
 import { PlusIcon } from '@heroicons/react/24/solid';
 import clsx from 'clsx';
 import {
-  OUTPUT_ID,
   ProcessingStep,
   ProcessingStepTypes,
   SourceStep,
@@ -47,7 +46,7 @@ export function NodeConnector({
   currentStep?: SourceStep | ProcessingStep;
   showAuto?: boolean;
 }) {
-  const litlytics = useLitlytics();
+  const { litlytics, setPipeline } = useLitlytics();
   const [step, setStep] = useState<ProcessingStep>(
     structuredClone(defaultStep)
   );
@@ -66,68 +65,13 @@ export function NodeConnector({
     }
 
     setLoading(true);
-    // generate new ID and double-check that it doesn't overlap with other steps
-    let id = litlytics.pipeline.steps.length;
-    let existingStep = litlytics.pipeline.steps.find(
-      (s) => s.id === `step_${id}`
-    );
-    while (existingStep) {
-      id += 1;
-      existingStep = litlytics.pipeline.steps.find(
-        (s) => s.id === `step_${id}`
-      );
-    }
-    // generate final id
-    const idStr = `step_${id}`;
-
-    let newStep: ProcessingStep | undefined = undefined;
-    if (manual) {
-      newStep = structuredClone(step);
-      newStep.id = idStr;
-    } else {
-      // generate new step
-      newStep = await litlytics.generateStep({
-        id: idStr,
-        name: step.name,
-        description: step.description,
-        input: step.input as StepInput,
-        type: step.type,
-      });
-    }
 
-    if (currentStep?.type === 'source') {
-      // connect new step to next node
-      const nextNodeId =
-        litlytics.pipeline.source.connectsTo.at(0) ?? OUTPUT_ID;
-      newStep.connectsTo = [nextNodeId];
-      // add
-      litlytics.setPipeline({
-        source: {
-          ...litlytics.pipeline.source,
-          connectsTo: [newStep.id],
-        },
-        steps: litlytics.pipeline.steps.concat(newStep),
-      });
-    } else {
-      // connect new step to next node
-      const nextNodeId =
-        litlytics.pipeline.steps
-          .find((s) => s.id === currentStep?.id)
-          ?.connectsTo.at(0) ?? OUTPUT_ID;
-      newStep.connectsTo = [nextNodeId];
-      // add
-      litlytics.setPipeline({
-        steps: litlytics.pipeline.steps
-          .map((s) => {
-            if (s.id === currentStep?.id) {
-              s.connectsTo = [newStep.id];
-              return s;
-            }
-            return s;
-          })
-          .concat(newStep),
-      });
-    }
+    const newPipeline = await litlytics.addStep({
+      step,
+      sourceStep: currentStep,
+      manual,
+    });
+    setPipeline(newPipeline);
 
     setIsOpen(false);
     setLoading(false);
diff --git a/app/components/pipeline/nodes/StepNode.tsx b/app/components/pipeline/nodes/StepNode.tsx
index c246c1b..0e2f967 100644
--- a/app/components/pipeline/nodes/StepNode.tsx
+++ b/app/components/pipeline/nodes/StepNode.tsx
@@ -36,7 +36,7 @@ import { useLitlytics } from '~/store/WithLitLytics';
 import { NodeContent, NodeFrame, NodeHeader } from './NodeFrame';
 
 export function StepNode({ data }: { data: ProcessingStep }) {
-  const litlytics = useLitlytics();
+  const { litlytics, pipeline, setPipeline, pipelineStatus } = useLitlytics();
   const [isOpen, setIsOpen] = useState(false);
   const [loading, setLoading] = useState(false);
   const [refine, setRefine] = useState('');
@@ -44,7 +44,7 @@ export function StepNode({ data }: { data: ProcessingStep }) {
   const { averageTiming, averagePrompt, averageCompletion, averageCost } =
     useMemo(() => {
       // const timings = data.
-      const results = litlytics.pipeline.resultDocs ?? [];
+      const results = pipeline.resultDocs ?? [];
       const res = results.filter((doc) => doc);
       if (!res.length) {
         return {};
@@ -73,19 +73,19 @@ export function StepNode({ data }: { data: ProcessingStep }) {
           completionTokens.length
       );
       const inputCost =
-        litlytics.config.provider === 'ollama'
+        litlytics.provider === 'ollama'
           ? 0
-          : modelCosts[litlytics.config.model!].input;
+          : modelCosts[litlytics.model!].input;
       const outputCost =
-        litlytics.config.provider === 'ollama'
+        litlytics.provider === 'ollama'
           ? 0
-          : modelCosts[litlytics.config.model!].output;
+          : modelCosts[litlytics.model!].output;
       const averageCost = _.round(
         averagePrompt * inputCost + averageCompletion * outputCost,
         3
       );
       return { averageTiming, averagePrompt, averageCompletion, averageCost };
-    }, [litlytics, data]);
+    }, [litlytics, pipeline, data]);
 
   const updateNodeByKey = (
     newVal: string | boolean | undefined,
@@ -94,13 +94,14 @@ export function StepNode({ data }: { data: ProcessingStep }) {
     const newData = structuredClone(data);
     newData[prop] = newVal;
 
-    const newSteps = litlytics.pipeline.steps.map((s) => {
+    const newSteps = pipeline.steps.map((s) => {
       if (s.id === newData.id) {
         return newData;
       }
       return s;
     });
-    litlytics.setPipeline({
+    setPipeline({
+      ...pipeline,
       steps: newSteps,
     });
   };
@@ -121,19 +122,20 @@ export function StepNode({ data }: { data: ProcessingStep }) {
     const newData = structuredClone(data);
     newData.code = newCode;
 
-    const newSteps = litlytics.pipeline.steps.map((s) => {
+    const newSteps = pipeline.steps.map((s) => {
       if (s.id === newData.id) {
         return newData;
       }
       return s;
     });
-    litlytics.setPipeline({
+    setPipeline({
+      ...pipeline,
       steps: newSteps,
     });
   };
 
   const deleteStep = () => {
-    const newSteps = litlytics.pipeline.steps
+    const newSteps = pipeline.steps
       // remove step
       .filter((s) => s.id !== data.id)
       // remove all links to step
@@ -147,16 +149,17 @@ export function StepNode({ data }: { data: ProcessingStep }) {
           connectsTo,
         };
       });
-    let sourceConnect = litlytics.pipeline.source.connectsTo.filter(
+    let sourceConnect = pipeline.source.connectsTo.filter(
       (id) => id !== data.id
     );
     const firstStep = newSteps.at(0);
     if (sourceConnect.length === 0 && firstStep) {
       sourceConnect = [firstStep.id];
     }
-    litlytics.setPipeline({
+    setPipeline({
+      ...pipeline,
       source: {
-        ...litlytics.pipeline.source,
+        ...pipeline.source,
         connectsTo: sourceConnect,
       },
       steps: newSteps,
@@ -175,13 +178,14 @@ export function StepNode({ data }: { data: ProcessingStep }) {
       refineRequest: refine,
       step: data,
     });
-    const newSteps = litlytics.pipeline.steps.map((s) => {
+    const newSteps = pipeline.steps.map((s) => {
       if (s.id === data.id) {
         return newStep;
       }
       return s;
     });
-    litlytics.setPipeline({
+    setPipeline({
+      ...pipeline,
       steps: newSteps,
     });
     setRefine('');
@@ -198,16 +202,16 @@ export function StepNode({ data }: { data: ProcessingStep }) {
         size={data.expanded ? 'sm' : 'collapsed'}
         className="pb-1"
         error={
-          litlytics.pipelineStatus.status === 'error' &&
-          litlytics.pipelineStatus.currentStep?.id === data.id
-            ? litlytics.pipelineStatus.error
+          pipelineStatus.status === 'error' &&
+          pipelineStatus.currentStep?.id === data.id
+            ? pipelineStatus.error
             : undefined
         }
       >
         <NodeHeader collapsed={!data.expanded}>
           <div className="flex flex-1 gap-2 items-center">
-            {litlytics.pipelineStatus.status === 'step' &&
-            litlytics.pipelineStatus.currentStep?.id === data.id ? (
+            {pipelineStatus.status === 'step' &&
+            pipelineStatus.currentStep?.id === data.id ? (
               <Spinner className="w-4 h-4" />
             ) : (
               <Button
diff --git a/app/components/pipeline/nodes/output/OutputNode.tsx b/app/components/pipeline/nodes/output/OutputNode.tsx
index 2fcf0d9..8dddf95 100644
--- a/app/components/pipeline/nodes/output/OutputNode.tsx
+++ b/app/components/pipeline/nodes/output/OutputNode.tsx
@@ -38,19 +38,22 @@ import { BasicOutputRender } from './BasicOutput';
 import { OutputType, OutputTypes } from './types';
 
 export function OutputNode() {
-  const litlytics = useLitlytics();
+  const {
+    litlytics,
+    pipeline,
+    setPipeline,
+    pipelineStatus,
+    setPipelineStatus,
+  } = useLitlytics();
   const [isOpen, setIsOpen] = useState(false);
   const [expanded, setExpanded] = useState(true);
   const [outputType, setOutputType] = useState<OutputType>('basic');
 
-  const results = useMemo(
-    () => litlytics.pipeline.results,
-    [litlytics.pipeline]
-  );
+  const results = useMemo(() => pipeline.results, [pipeline]);
 
   const { timing, prompt, completion, cost } = useMemo(() => {
     // const timings = data.
-    const results = litlytics.pipeline.resultDocs ?? [];
+    const results = pipeline.resultDocs ?? [];
     const res = results.filter((doc) => doc);
     if (!res.length) {
       return {};
@@ -68,33 +71,25 @@ export function OutputNode() {
     const prompt = promptTokens.reduce((acc, val) => acc + val, 0);
     const completion = completionTokens.reduce((acc, val) => acc + val, 0);
     const inputCost =
-      litlytics.config.provider === 'ollama'
-        ? 0
-        : modelCosts[litlytics.config.model!].input;
+      litlytics.provider === 'ollama' ? 0 : modelCosts[litlytics.model!].input;
     const outputCost =
-      litlytics.config.provider === 'ollama'
-        ? 0
-        : modelCosts[litlytics.config.model!].output;
+      litlytics.provider === 'ollama' ? 0 : modelCosts[litlytics.model!].output;
     const cost = _.round(prompt * inputCost + completion * outputCost, 3);
     return { timing, prompt, completion, cost };
-  }, [litlytics]);
+  }, [litlytics, pipeline]);
 
   const doRunPipeline = async () => {
-    await litlytics.runPipeline({
-      // very hacky, but will trigger UI re-renders
-      // while setting status internally won't because proxy doesn't see those triggers
-      onStatus: (status) => {
-        litlytics.setPipelineStatus(status);
-      },
+    const newPipeline = await litlytics.runPipeline({
+      onStatus: (status) => setPipelineStatus(status),
     });
+    setPipeline(newPipeline);
 
     // write final status to update on end
-    litlytics.setPipelineStatus({ status: 'done' });
+    setPipelineStatus({ status: 'done' });
   };
 
   const running =
-    litlytics.pipelineStatus.status === 'sourcing' ||
-    litlytics.pipelineStatus.status === 'step';
+    pipelineStatus.status === 'sourcing' || pipelineStatus.status === 'step';
 
   return (
     <>
@@ -122,7 +117,7 @@ export function OutputNode() {
                 <PlayIcon aria-hidden="true" className="h-5 w-5" />
               )}
             </Button>
-            {litlytics.pipelineStatus.status === 'done' && (
+            {pipelineStatus.status === 'done' && (
               <div
                 className="flex flex-1 justify-end"
                 title="Successfully finished!"
@@ -130,7 +125,7 @@ export function OutputNode() {
                 <CheckIcon color="green" className="w-5 h-5" />
               </div>
             )}
-            {litlytics.pipelineStatus.status === 'error' && (
+            {pipelineStatus.status === 'error' && (
               <div
                 className="flex flex-1 justify-end"
                 title="Error during execution!"
@@ -182,7 +177,7 @@ export function OutputNode() {
                 </Badge>
               )}
             </div>
-            <BasicOutputRender pipeline={litlytics.pipeline} />
+            <BasicOutputRender pipeline={pipeline} />
           </NodeContent>
         ) : (
           <></>
diff --git a/app/components/pipeline/nodes/source/SourceNode.tsx b/app/components/pipeline/nodes/source/SourceNode.tsx
index d2e8940..f95a40d 100644
--- a/app/components/pipeline/nodes/source/SourceNode.tsx
+++ b/app/components/pipeline/nodes/source/SourceNode.tsx
@@ -42,15 +42,12 @@ const UnknownSource = ({
 
 export function SourceNode() {
   const [isOpen, setIsOpen] = useState(false);
-  const litlytics = useLitlytics();
+  const { litlytics, pipeline, setPipeline, pipelineStatus } = useLitlytics();
 
   const source = useMemo(() => {
-    return litlytics.pipeline.source;
-  }, [litlytics.pipeline.source]);
-  const docs = useMemo(
-    () => litlytics.pipeline.source.docs,
-    [litlytics.pipeline.source.docs]
-  );
+    return pipeline.source;
+  }, [pipeline.source]);
+  const docs = useMemo(() => source.docs, [source.docs]);
 
   const sourceType = useMemo(
     () => ((source.config as SourceConfig)?.type ?? 'text') as SourceType,
@@ -72,13 +69,15 @@ export function SourceNode() {
   ) => {
     const newData = structuredClone(source!);
     newData[prop] = newVal;
-    litlytics.setPipeline({
+    setPipeline({
+      ...pipeline,
       source: newData,
     });
   };
 
   const updateDocs = (newDocs: Doc[]) => {
-    litlytics.setDocs(newDocs);
+    const newPipeline = litlytics.setDocs(newDocs);
+    setPipeline(newPipeline);
   };
 
   if (!source) {
@@ -89,7 +88,7 @@ export function SourceNode() {
     <>
       {/* Source node render */}
       <NodeFrame
-        hasConnector={litlytics.pipeline.steps.length > 0 ? true : 'auto'}
+        hasConnector={pipeline.steps.length > 0 ? true : 'auto'}
         currentStep={source}
         size={
           // always collapse empty
@@ -104,7 +103,7 @@ export function SourceNode() {
           collapsed={sourceType === 'empty' ? true : !source.expanded}
         >
           <div className="flex flex-1 gap-2 items-center">
-            {litlytics.pipelineStatus.status === 'sourcing' ? (
+            {pipelineStatus.status === 'sourcing' ? (
               <Spinner className="w-4 h-4" />
             ) : (
               <Button
@@ -159,9 +158,10 @@ export function SourceNode() {
                 name="step-input"
                 value={sourceType as string}
                 onChange={(e) =>
-                  litlytics.setPipeline({
+                  setPipeline({
+                    ...pipeline,
                     source: {
-                      ...litlytics.pipeline.source,
+                      ...pipeline.source,
                       config: {
                         type: e.target.value as SourceType,
                       } as SourceConfig,
diff --git a/app/components/pipeline/source/useTestDocs.ts b/app/components/pipeline/source/useTestDocs.ts
index 72c3623..3c2fc7c 100644
--- a/app/components/pipeline/source/useTestDocs.ts
+++ b/app/components/pipeline/source/useTestDocs.ts
@@ -2,11 +2,8 @@ import { useMemo } from 'react';
 import { useLitlytics } from '~/store/WithLitLytics';
 
 export function useTestDocs() {
-  const litlytics = useLitlytics();
-  const allDocs = useMemo(
-    () => litlytics.pipeline.source.docs,
-    [litlytics.pipeline.source.docs]
-  );
+  const { pipeline } = useLitlytics();
+  const allDocs = useMemo(() => pipeline.source.docs, [pipeline.source.docs]);
   const testDocs = useMemo(() => allDocs.filter((d) => d?.test), [allDocs]);
 
   return { allDocs, testDocs };
diff --git a/app/components/step/StepTest.tsx b/app/components/step/StepTest.tsx
index 9fc968d..62a5cfb 100644
--- a/app/components/step/StepTest.tsx
+++ b/app/components/step/StepTest.tsx
@@ -17,7 +17,7 @@ import { CustomMarkdown } from '../markdown/Markdown';
 import { useTestDocs } from '../pipeline/source/useTestDocs';
 
 export function StepTest({ data }: { data: ProcessingStep }) {
-  const litlytics = useLitlytics();
+  const { litlytics, setPipeline } = useLitlytics();
   const { testDocs, allDocs } = useTestDocs();
   const [testDocId, setTestDocId] = useState(testDocs.at(0)?.id ?? '');
   const [isTestOpen, setTestOpen] = useState(false);
@@ -65,7 +65,8 @@ export function StepTest({ data }: { data: ProcessingStep }) {
           }
           return d;
         });
-        litlytics.setDocs(newDocs);
+        const newPipeline = litlytics.setDocs(newDocs);
+        setPipeline(newPipeline);
       } else {
         const newDocs = allDocs.map((d) => {
           if (d.id === doc?.id) {
@@ -73,7 +74,8 @@ export function StepTest({ data }: { data: ProcessingStep }) {
           }
           return d;
         });
-        litlytics.setDocs(newDocs);
+        const newPipeline = litlytics.setDocs(newDocs);
+        setPipeline(newPipeline);
       }
     } catch (err) {
       setError(err as Error);
diff --git a/app/components/ui/Overlay.tsx b/app/components/ui/Overlay.tsx
index 0a559b9..beb5d88 100644
--- a/app/components/ui/Overlay.tsx
+++ b/app/components/ui/Overlay.tsx
@@ -29,7 +29,7 @@ import {
   DropdownMenu,
   DropdownShortcut,
 } from '~/components/catalyst/dropdown';
-import { configAtom, configUndoAtom, webllmAtom } from '~/store/store';
+import { configAtom, pipelineUndoAtom, webllmAtom } from '~/store/store';
 import { useLitlytics } from '~/store/WithLitLytics';
 import { Field, FieldGroup, Label } from '../catalyst/fieldset';
 import { Input } from '../catalyst/input';
@@ -62,10 +62,11 @@ function MenuHolder({
 export function OverlayUI() {
   const fileInputRef = useRef<HTMLInputElement>(null);
   const loadInputRef = useRef<HTMLTextAreaElement>(null);
-  const litlytics = useLitlytics();
+  const { litlytics, pipeline, setPipeline, setPipelineStatus } =
+    useLitlytics();
   const litlyticsConfig = useAtomValue(configAtom);
   const webllm = useAtomValue(webllmAtom);
-  const { undo, redo, canUndo, canRedo } = useAtomValue(configUndoAtom);
+  const { undo, redo, canUndo, canRedo } = useAtomValue(pipelineUndoAtom);
   const [isOpen, setIsOpen] = useState(false);
   const [isSaveOpen, setIsSaveOpen] = useState(false);
   const [isLoadOpen, setIsLoadOpen] = useState(false);
@@ -95,12 +96,13 @@ export function OverlayUI() {
   };
 
   const resetPipeline = () => {
-    litlytics.resetPipeline();
+    const newPipeline = litlytics.resetPipeline();
+    setPipeline(newPipeline);
     setIsOpen(false);
   };
 
   const pipelineToJSON = () => {
-    const pipelineToSave = structuredClone(litlytics.pipeline);
+    const pipelineToSave = structuredClone(pipeline);
     // set model and provider
     pipelineToSave.provider = litlyticsConfig.provider;
     pipelineToSave.model = litlyticsConfig.model;
@@ -119,7 +121,7 @@ export function OverlayUI() {
     // Create a temporary anchor element
     const a = document.createElement('a');
     a.href = url;
-    a.download = `${litlytics.pipeline.name}.json`;
+    a.download = `${pipeline.name}.json`;
     // Programmatically click the anchor to trigger the download
     a.click();
     // Clean up: revoke the object URL
@@ -151,8 +153,8 @@ export function OverlayUI() {
     }
     try {
       const pipelineJson = JSON.parse(inputStr);
-      litlytics.setPipeline(pipelineJson);
-      litlytics.setPipelineStatus({ status: 'init' });
+      setPipeline(pipelineJson);
+      setPipelineStatus({ status: 'init' });
       loadInputRef.current!.value = '';
       setIsLoadOpen(false);
     } catch (err) {
@@ -167,8 +169,8 @@ export function OverlayUI() {
       reader.onload = (e) => {
         try {
           const json = JSON.parse(e.target?.result as string) as Pipeline;
-          litlytics.setPipeline(json);
-          litlytics.setPipelineStatus({ status: 'init' });
+          setPipeline(json);
+          setPipelineStatus({ status: 'init' });
           setIsLoadOpen(false);
         } catch (error) {
           console.error('Error parsing JSON:', error);
@@ -337,8 +339,8 @@ export function OverlayUI() {
                 name="name"
                 placeholder="Pipeline name"
                 autoFocus
-                value={litlytics.pipeline.name}
-                onChange={(e) => (litlytics.pipeline.name = e.target.value)}
+                value={pipeline.name}
+                onChange={(e) => (pipeline.name = e.target.value)}
               />
             </Field>
           </FieldGroup>
diff --git a/app/store/WithLitLytics.tsx b/app/store/WithLitLytics.tsx
index 1e09746..083383a 100644
--- a/app/store/WithLitLytics.tsx
+++ b/app/store/WithLitLytics.tsx
@@ -1,63 +1,38 @@
 import { useAtom, useAtomValue } from 'jotai';
-import { LitLytics } from 'litlytics';
-import {
-  createContext,
-  useContext,
-  useEffect,
-  useReducer,
-  useRef,
-} from 'react';
+import { emptyPipeline, LitLytics, Pipeline, PipelineStatus } from 'litlytics';
+import { createContext, useContext, useEffect } from 'react';
 import {
   configAtom,
   litlyticsAtom,
-  litlyticsInitedAtom,
+  pipelineAtom,
+  pipelineStatusAtom,
   webllmAtom,
 } from './store';
-import { createReactiveProxy } from './util';
 
-const LitLyticsContext = createContext<{ litlytics?: LitLytics }>({
-  litlytics: undefined,
+const LitLyticsContext = createContext<{
+  litlytics: LitLytics;
+  pipeline: Pipeline;
+  setPipeline: (p: Pipeline) => void;
+  pipelineStatus: PipelineStatus;
+  setPipelineStatus: (s: PipelineStatus) => void;
+}>({
+  litlytics: new LitLytics({
+    provider: 'openai',
+    model: 'gpt-4o-mini',
+    key: '',
+  }),
+  pipeline: emptyPipeline,
+  setPipeline: () => {},
+  pipelineStatus: { status: 'init' } as PipelineStatus,
+  setPipelineStatus: () => {},
 });
 
 export function WithLitLytics({ children }: { children: React.ReactNode }) {
-  const saveConfigRef = useRef<Timer>();
   const webllm = useAtomValue(webllmAtom);
-  const [config, setConfig] = useAtom(configAtom);
-  const [isInited, setIsInited] = useAtom(litlyticsInitedAtom);
-  const [litlytics, setLitlytics] = useAtom(litlyticsAtom);
-  const [, forceUpdate] = useReducer((x) => x + 1, 0);
-
-  // create proxy that triggers state updates
-  // only create it once and use context to pass it around
-  useEffect(() => {
-    if (isInited) {
-      return;
-    }
-    const ll = new LitLytics({
-      provider: 'openai',
-      model: 'gpt-4o-mini',
-      key: '',
-    });
-    setLitlytics(
-      createReactiveProxy(ll, () => {
-        // if there's current save timeout - reset it
-        if (saveConfigRef.current) {
-          clearTimeout(saveConfigRef.current);
-        }
-        // persist changes in localStorage after some time of inactivity
-        saveConfigRef.current = setTimeout(() => {
-          setConfig(ll.exportConfig());
-        }, 500);
-        // force UI update
-        forceUpdate();
-      })
-    );
-    // set inited flag
-    setIsInited(true);
-    // assign initial config and webEngine instance
-    ll.importConfig(config);
-    ll.setWebEngine(webllm.engine);
-  }, [isInited, config, webllm, setIsInited, setLitlytics, setConfig]);
+  const config = useAtomValue(configAtom);
+  const litlytics = useAtomValue(litlyticsAtom);
+  const [pipeline, setPipeline] = useAtom(pipelineAtom);
+  const [pipelineStatus, setPipelineStatus] = useAtom(pipelineStatusAtom);
 
   // update config / llm on changes
   useEffect(() => {
@@ -71,19 +46,32 @@ export function WithLitLytics({ children }: { children: React.ReactNode }) {
       litlytics.importConfig(config);
     }
     // assign webllm engine
-    if (litlytics.getEngine() !== webllm.engine) {
-      litlytics.setWebEngine(webllm.engine);
+    if (litlytics.engine !== webllm.engine) {
+      litlytics.engine = webllm.engine;
     }
   }, [config, webllm, litlytics]);
 
+  // update pipeline on changes
+  useEffect(() => {
+    litlytics.setPipeline(pipeline);
+  }, [pipeline, litlytics]);
+
   return (
-    <LitLyticsContext.Provider value={{ litlytics }}>
+    <LitLyticsContext.Provider
+      value={{
+        litlytics,
+        pipeline,
+        setPipeline,
+        pipelineStatus,
+        setPipelineStatus,
+      }}
+    >
       {children}
     </LitLyticsContext.Provider>
   );
 }
 
 export const useLitlytics = () => {
-  const { litlytics } = useContext(LitLyticsContext);
-  return litlytics!;
+  const ctx = useContext(LitLyticsContext);
+  return ctx;
 };
diff --git a/app/store/store.ts b/app/store/store.ts
index bbef491..db82406 100644
--- a/app/store/store.ts
+++ b/app/store/store.ts
@@ -1,8 +1,16 @@
 import { atom } from 'jotai';
 import { withUndo } from 'jotai-history';
 import { atomWithStorage } from 'jotai/utils';
-import { LitLytics, LitLyticsConfig, MLCEngine } from 'litlytics';
+import {
+  emptyPipeline,
+  LitLytics,
+  LitLyticsConfig,
+  MLCEngine,
+  Pipeline,
+  PipelineStatus,
+} from 'litlytics';
 
+// litlytics config, persisted in localStorage
 export const configAtom = atomWithStorage<LitLyticsConfig>(
   'litlytics.config',
   {
@@ -14,8 +22,20 @@ export const configAtom = atomWithStorage<LitLyticsConfig>(
   undefined,
   { getOnInit: true }
 );
-export const configUndoAtom = withUndo(configAtom, 10);
 
+// litlytics pipeline, persisted in localStorage
+export const pipelineAtom = atomWithStorage<Pipeline>(
+  'litlytics.pipeline',
+  emptyPipeline,
+  undefined,
+  { getOnInit: true }
+);
+export const pipelineUndoAtom = withUndo(pipelineAtom, 10);
+
+// pipeline status
+export const pipelineStatusAtom = atom<PipelineStatus>({ status: 'init' });
+
+// litlytics instance
 export const litlyticsAtom = atom<LitLytics>(
   new LitLytics({
     provider: 'openai',
@@ -24,6 +44,7 @@ export const litlyticsAtom = atom<LitLytics>(
   })
 );
 
+// webllm instance
 export const webllmAtom = atom<{
   engine?: MLCEngine;
   fetchProgress: number;
@@ -35,5 +56,3 @@ export const webllmAtom = atom<{
   loadProgress: -1,
   status: '',
 });
-
-export const litlyticsInitedAtom = atom(false);
diff --git a/app/store/util.ts b/app/store/util.ts
deleted file mode 100644
index c9c15a6..0000000
--- a/app/store/util.ts
+++ /dev/null
@@ -1,36 +0,0 @@
-import { LitLytics } from 'litlytics';
-/* eslint @typescript-eslint/ban-ts-comment: off */
-
-/**
- * Creates a Proxy wrapper around LitLytics instance to monitor function invocations
- * and update react state upon those invocations completion.
- * This is required because React cannot see changes in persistent object (i.e. wants new object refs)
- *
- * @param instance LitLytics instance
- * @param updateCallback callback to update state
- * @returns LitLytics instance with a proxy
- */
-export function createReactiveProxy(
-  instance: LitLytics,
-  updateCallback: () => void
-) {
-  return new Proxy(instance, {
-    get(target, prop, receiver) {
-      const value = Reflect.get(target, prop, receiver);
-      if (typeof value === 'function') {
-        // @ts-ignore
-        return function (...args) {
-          // @ts-ignore
-          const result = value.apply(this, args);
-          if (result instanceof Promise) {
-            result.then(() => updateCallback.call(target));
-          } else {
-            updateCallback.call(target);
-          }
-          return result;
-        };
-      }
-      return value;
-    },
-  });
-}
diff --git a/packages/litlytics/litlytics.ts b/packages/litlytics/litlytics.ts
index df8e4db..634f691 100644
--- a/packages/litlytics/litlytics.ts
+++ b/packages/litlytics/litlytics.ts
@@ -26,7 +26,11 @@ import { refinePipeline } from './pipeline/refine';
 import { generateCodeExplain } from './step/explain';
 import { generateStep, type GenerateStepArgs } from './step/generate';
 import { refineStep } from './step/refine';
-import { type ProcessingStep } from './step/Step';
+import {
+  type ProcessingStep,
+  type SourceStep,
+  type StepInput,
+} from './step/Step';
 
 // export types and commonly used vars
 export type { MLCEngine } from '@mlc-ai/web-llm';
@@ -40,7 +44,11 @@ export {
   type LLMProvider,
 } from './llm/types';
 export { OUTPUT_ID } from './output/Output';
-export type { Pipeline, PipelineStatus } from './pipeline/Pipeline';
+export {
+  emptyPipeline,
+  type Pipeline,
+  type PipelineStatus,
+} from './pipeline/Pipeline';
 export {
   StepInputs,
   type ProcessingStep,
@@ -214,7 +222,7 @@ export class LitLytics {
     onStatus: ({ step, totalSteps }: PipelineFromTextStatus) => void
   ) => {
     if (!this.pipeline.pipelinePlan) {
-      return;
+      return this.pipeline;
     }
 
     const newSteps = await pipelineFromText(
@@ -240,7 +248,7 @@ export class LitLytics {
 
   generatePipeline = async () => {
     if (!this.pipeline.pipelineDescription?.length) {
-      return;
+      return this.pipeline;
     }
 
     const plan = await generatePipeline({
@@ -273,13 +281,9 @@ export class LitLytics {
       this.setPipelineStatus(status);
       onStatus?.(this.pipelineStatus);
     };
-    try {
-      setStatus({ status: 'init' });
-      const newPipeline = await runPipeline(this, setStatus);
-      return this.setPipeline(newPipeline);
-    } catch (err) {
-      setStatus({ status: 'error', error: err as Error });
-    }
+    setStatus({ status: 'init' });
+    const newPipeline = await runPipeline(this, setStatus);
+    return this.setPipeline(newPipeline);
   };
 
   /**
@@ -306,6 +310,74 @@ export class LitLytics {
     return await generateCodeExplain({ litlytics: this, code });
   };
 
+  addStep = async ({
+    step,
+    sourceStep,
+    manual,
+  }: {
+    step: ProcessingStep;
+    sourceStep?: SourceStep | ProcessingStep;
+    manual?: boolean;
+  }) => {
+    // generate new ID and double-check that it doesn't overlap with other steps
+    let id = this.pipeline.steps.length;
+    let existingStep = this.pipeline.steps.find((s) => s.id === `step_${id}`);
+    while (existingStep) {
+      id += 1;
+      existingStep = this.pipeline.steps.find((s) => s.id === `step_${id}`);
+    }
+    // generate final id
+    const idStr = `step_${id}`;
+
+    let newStep: ProcessingStep | undefined = undefined;
+    if (manual) {
+      newStep = structuredClone(step);
+      newStep.id = idStr;
+    } else {
+      // generate new step
+      newStep = await this.generateStep({
+        id: idStr,
+        name: step.name,
+        description: step.description,
+        input: step.input as StepInput,
+        type: step.type,
+      });
+    }
+
+    if (sourceStep?.type === 'source') {
+      // connect new step to next node
+      const nextNodeId = this.pipeline.source.connectsTo.at(0) ?? OUTPUT_ID;
+      newStep.connectsTo = [nextNodeId];
+      // add
+      return this.setPipeline({
+        source: {
+          ...this.pipeline.source,
+          connectsTo: [newStep.id],
+        },
+        steps: this.pipeline.steps.concat(newStep),
+      });
+    }
+
+    // connect new step to next node
+    const nextNodeId =
+      this.pipeline.steps
+        .find((s) => s.id === sourceStep?.id)
+        ?.connectsTo.at(0) ?? OUTPUT_ID;
+    newStep.connectsTo = [nextNodeId];
+    // add
+    return this.setPipeline({
+      steps: this.pipeline.steps
+        .map((s) => {
+          if (s.id === sourceStep?.id) {
+            s.connectsTo = [newStep.id];
+            return s;
+          }
+          return s;
+        })
+        .concat(newStep),
+    });
+  };
+
   refineStep = async ({
     refineRequest,
     step,
diff --git a/packages/litlytics/package.json b/packages/litlytics/package.json
index d636003..3f276c9 100644
--- a/packages/litlytics/package.json
+++ b/packages/litlytics/package.json
@@ -4,7 +4,7 @@
   "main": "dist/litlytics.js",
   "types": "litlytics.ts",
   "type": "module",
-  "version": "0.3.1",
+  "version": "0.3.2",
   "scripts": {
     "build": "bun run build.ts",
     "typecheck": "tsc",
diff --git a/packages/litlytics/test/litlytics.test.ts b/packages/litlytics/test/litlytics.test.ts
index fbaa197..8775ec1 100644
--- a/packages/litlytics/test/litlytics.test.ts
+++ b/packages/litlytics/test/litlytics.test.ts
@@ -1,6 +1,11 @@
 import type { LanguageModelUsage } from 'ai';
 import { expect, test } from 'vitest';
-import { LitLytics, type Pipeline } from '../litlytics';
+import {
+  LitLytics,
+  OUTPUT_ID,
+  type Pipeline,
+  type StepInput,
+} from '../litlytics';
 import { docSecond, firstStep, pipeline } from './data/testPipeline';
 
 test('should create new instance and import/export config', async () => {
@@ -60,3 +65,126 @@ test('should create new instance and run test pipeline', async () => {
   expect(newPipeline?.results?.[0].result).toEqual('test doc result second');
   expect(newPipeline?.results?.[1].result).toEqual('second doc result second');
 });
+
+test('should add new step to pipeline after source', async () => {
+  let userPrompt = '';
+  const litlytics = new LitLytics({
+    provider: 'openai',
+    model: 'test',
+    key: 'test',
+  });
+  litlytics.setPipeline(pipeline);
+  // mock prompt replies
+  litlytics.runPrompt = async ({ user }: { system: string; user: string }) => {
+    userPrompt = user;
+    return { result: `New step prompt`, usage: {} as LanguageModelUsage };
+  };
+  // execute pipeline
+  const newPipeline = await litlytics.addStep({
+    step: {
+      id: '0',
+      name: 'New step',
+      description: 'New step',
+      type: 'llm',
+      input: 'doc' as StepInput,
+      connectsTo: [pipeline.steps.at(0)!.id],
+    },
+    sourceStep: pipeline.source,
+  });
+  expect(userPrompt).toEqual(`Step name: New step
+Step description: New step
+Step input: doc`);
+  // validate pipeline
+  expect(newPipeline).toBeDefined();
+  expect(newPipeline?.steps).toHaveLength(3);
+  // validate step
+  const newStep = newPipeline?.steps?.find((s) => s.name === 'New step');
+  expect(newStep?.id).toEqual('step_2');
+  expect(newStep?.name).toEqual('New step');
+  expect(newStep?.prompt).toEqual('New step prompt');
+  // validate connections
+  expect(newPipeline.source.connectsTo).toEqual([newStep!.id]);
+  expect(newStep?.connectsTo).toEqual([newPipeline.steps.at(0)!.id]);
+});
+
+test('should add new step to pipeline after step', async () => {
+  let userPrompt = '';
+  const litlytics = new LitLytics({
+    provider: 'openai',
+    model: 'test',
+    key: 'test',
+  });
+  litlytics.setPipeline(pipeline);
+  // mock prompt replies
+  litlytics.runPrompt = async ({ user }: { system: string; user: string }) => {
+    userPrompt = user;
+    return { result: `New step prompt`, usage: {} as LanguageModelUsage };
+  };
+  // execute pipeline
+  const newPipeline = await litlytics.addStep({
+    step: {
+      id: '0',
+      name: 'New step',
+      description: 'New step',
+      type: 'llm',
+      input: 'doc' as StepInput,
+      connectsTo: [pipeline.steps.at(0)!.id],
+    },
+    sourceStep: pipeline.steps.at(0),
+  });
+  expect(userPrompt).toEqual(`Step name: New step
+Step description: New step
+Step input: doc`);
+  // validate pipeline
+  expect(newPipeline).toBeDefined();
+  expect(newPipeline?.steps).toHaveLength(3);
+  // validate step
+  const newStep = newPipeline?.steps?.find((s) => s.name === 'New step');
+  expect(newStep?.id).toEqual('step_2');
+  expect(newStep?.name).toEqual('New step');
+  expect(newStep?.prompt).toEqual('New step prompt');
+  // validate connections
+  expect(newPipeline.steps.at(0)!.connectsTo).toEqual([newStep!.id]);
+  expect(newStep?.connectsTo).toEqual([newPipeline.steps.at(1)!.id]);
+});
+
+test('should add new step to pipeline at the end', async () => {
+  let userPrompt = '';
+  const litlytics = new LitLytics({
+    provider: 'openai',
+    model: 'test',
+    key: 'test',
+  });
+  litlytics.setPipeline(pipeline);
+  // mock prompt replies
+  litlytics.runPrompt = async ({ user }: { system: string; user: string }) => {
+    userPrompt = user;
+    return { result: `New step prompt`, usage: {} as LanguageModelUsage };
+  };
+  // execute pipeline
+  const newPipeline = await litlytics.addStep({
+    step: {
+      id: '0',
+      name: 'New step',
+      description: 'New step',
+      type: 'llm',
+      input: 'doc' as StepInput,
+      connectsTo: [],
+    },
+    sourceStep: pipeline.steps.at(-1),
+  });
+  expect(userPrompt).toEqual(`Step name: New step
+Step description: New step
+Step input: doc`);
+  // validate pipeline
+  expect(newPipeline).toBeDefined();
+  expect(newPipeline?.steps).toHaveLength(3);
+  // validate step
+  const newStep = newPipeline?.steps?.find((s) => s.name === 'New step');
+  expect(newStep?.id).toEqual('step_2');
+  expect(newStep?.name).toEqual('New step');
+  expect(newStep?.prompt).toEqual('New step prompt');
+  // validate connections
+  expect(newPipeline.steps.at(1)!.connectsTo).toEqual([newStep!.id]);
+  expect(newStep?.connectsTo).toEqual([OUTPUT_ID]);
+});