File size: 2,383 Bytes
cd6f98e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
import type { StateCreator } from "zustand";
import { create } from "zustand";
import { createJSONStorage, persist } from "zustand/middleware";

import { createSelectors } from "./helpers";
import type { ActiveTool } from "../hooks/useTools";
import type { AgentLifecycle } from "../services/agent/agent-run-model";
import type AutonomousAgent from "../services/agent/autonomous-agent";

interface AgentSlice {
  agent: AutonomousAgent | null;
  lifecycle: AgentLifecycle;
  setLifecycle: (AgentLifecycle) => void;
  summarized: boolean;
  setSummarized: (boolean) => void;
  isAgentThinking: boolean;
  setIsAgentThinking: (isThinking: boolean) => void;
  setAgent: (newAgent: AutonomousAgent | null) => void;
}

const initialAgentState = {
  agent: null,
  lifecycle: "offline" as const,
  summarized: false,
  isAgentThinking: false,
  isAgentPaused: undefined,
};

interface ToolsSlice {
  tools: Omit<ActiveTool, "active">[];
  setTools: (tools: ActiveTool[]) => void;
}

const resetters: (() => void)[] = [];

const createAgentSlice: StateCreator<AgentSlice> = (set, get) => {
  resetters.push(() => set(initialAgentState));
  return {
    ...initialAgentState,
    setLifecycle: (lifecycle: AgentLifecycle) => {
      set(() => ({
        lifecycle: lifecycle,
      }));
    },
    setSummarized: (summarized: boolean) => {
      set(() => ({
        summarized: summarized,
      }));
    },
    setIsAgentThinking: (isThinking: boolean) => {
      set(() => ({
        isAgentThinking: isThinking,
      }));
    },
    setAgent: (newAgent) => {
      set(() => ({
        agent: newAgent,
      }));

      if (get().agent === null) {
        resetters.forEach((resetter) => resetter());
      }
    },
  };
};

const createToolsSlice: StateCreator<ToolsSlice> = (set) => {
  return {
    tools: [],
    setTools: (tools) => {
      set(() => ({
        tools: tools,
      }));
    },
  };
};

export const useAgentStore = createSelectors(
  create<AgentSlice & ToolsSlice>()(
    persist(
      (...a) => ({
        ...createAgentSlice(...a),
        ...createToolsSlice(...a),
      }),
      {
        name: "agent-storage-v2",
        storage: createJSONStorage(() => localStorage),
        partialize: (state) => ({
          tools: state.tools,
        }),
      }
    )
  )
);

export const resetAllAgentSlices = () => resetters.forEach((resetter) => resetter());