agent-swarm-kit 1.0.68 → 1.0.70

package/build/index.cjs

@@ -236,6 +236,7 @@ var connectionServices$1 = {
     swarmConnectionService: Symbol('swarmConnectionService'),
     sessionConnectionService: Symbol('sessionConnectionService'),
     storageConnectionService: Symbol('storageConnectionService'),
+    stateConnectionService: Symbol('stateConnectionService'),
 };
 var schemaServices$1 = {
     completionSchemaService: Symbol('completionSchemaService'),
@@ -244,6 +245,7 @@ var schemaServices$1 = {
     toolSchemaService: Symbol('toolSchemaService'),
     embeddingSchemaService: Symbol('embeddingSchemaService'),
     storageSchemaService: Symbol('storageSchemaService'),
+    stateSchemaService: Symbol('stateSchemaService'),
 };
 var publicServices$1 = {
     agentPublicService: Symbol('agentPublicService'),
@@ -251,6 +253,7 @@ var publicServices$1 = {
     sessionPublicService: Symbol('sessionPublicService'),
     swarmPublicService: Symbol('swarmPublicService'),
     storagePublicService: Symbol('storagePublicService'),
+    statePublicService: Symbol('statePublicService'),
 };
 var validationServices$1 = {
     agentValidationService: Symbol('agentValidationService'),
@@ -1425,7 +1428,7 @@ var AgentConnectionService = /** @class */ (function () {
             var _b = __read(_a, 2), clientId = _b[0], agentName = _b[1];
             return "".concat(clientId, "-").concat(agentName);
         }, function (clientId, agentName) {
-            var _a = _this.agentSchemaService.get(agentName), prompt = _a.prompt, system = _a.system, tools = _a.tools, _b = _a.transform, transform = _b === void 0 ? GLOBAL_CONFIG.CC_AGENT_OUTPUT_TRANSFORM : _b, _c = _a.map, map = _c === void 0 ? GLOBAL_CONFIG.CC_AGENT_OUTPUT_MAP : _c, callbacks = _a.callbacks, storages = _a.storages, completionName = _a.completion, _d = _a.validate, validate = _d === void 0 ? validateDefault : _d;
+            var _a = _this.agentSchemaService.get(agentName), prompt = _a.prompt, system = _a.system, tools = _a.tools, _b = _a.transform, transform = _b === void 0 ? GLOBAL_CONFIG.CC_AGENT_OUTPUT_TRANSFORM : _b, _c = _a.map, map = _c === void 0 ? GLOBAL_CONFIG.CC_AGENT_OUTPUT_MAP : _c, callbacks = _a.callbacks, storages = _a.storages, states = _a.states, completionName = _a.completion, _d = _a.validate, validate = _d === void 0 ? validateDefault : _d;
             var completion = _this.completionSchemaService.get(completionName);
             _this.sessionValidationService.addAgentUsage(clientId, agentName);
             storages === null || storages === void 0 ? void 0 : storages.forEach(function (storageName) {
@@ -1433,6 +1436,11 @@ var AgentConnectionService = /** @class */ (function () {
                     .getStorage(clientId, storageName)
                     .waitForInit();
             });
+            states === null || states === void 0 ? void 0 : states.forEach(function (storageName) {
+                return _this.storageConnectionService
+                    .getStorage(clientId, storageName)
+                    .waitForInit();
+            });
             return new ClientAgent(__assign({ clientId: clientId, agentName: agentName, validate: validate, logger: _this.loggerService, history: _this.historyConnectionService.getHistory(clientId, agentName), prompt: prompt, system: system, transform: transform, map: map, tools: tools === null || tools === void 0 ? void 0 : tools.map(_this.toolSchemaService.get), completion: completion }, callbacks));
         });
         /**
@@ -2783,6 +2791,7 @@ var AgentPublicService = /** @class */ (function () {
                                 agentName: agentName,
                                 swarmName: "",
                                 storageName: "",
+                                stateName: "",
                             })];
                     case 1: return [2 /*return*/, _a.sent()];
                 }
@@ -2818,6 +2827,7 @@ var AgentPublicService = /** @class */ (function () {
                                 agentName: agentName,
                                 swarmName: "",
                                 storageName: "",
+                                stateName: "",
                             })];
                     case 1: return [2 /*return*/, _a.sent()];
                 }
@@ -2850,6 +2860,7 @@ var AgentPublicService = /** @class */ (function () {
                                 agentName: agentName,
                                 swarmName: "",
                                 storageName: "",
+                                stateName: "",
                             })];
                     case 1: return [2 /*return*/, _a.sent()];
                 }
@@ -2886,6 +2897,7 @@ var AgentPublicService = /** @class */ (function () {
                                 agentName: agentName,
                                 swarmName: "",
                                 storageName: "",
+                                stateName: "",
                             })];
                     case 1: return [2 /*return*/, _a.sent()];
                 }
@@ -2920,6 +2932,7 @@ var AgentPublicService = /** @class */ (function () {
                                 agentName: agentName,
                                 swarmName: "",
                                 storageName: "",
+                                stateName: "",
                             })];
                     case 1: return [2 /*return*/, _a.sent()];
                 }
@@ -2954,6 +2967,7 @@ var AgentPublicService = /** @class */ (function () {
                                 agentName: agentName,
                                 swarmName: "",
                                 storageName: "",
+                                stateName: "",
                             })];
                     case 1: return [2 /*return*/, _a.sent()];
                 }
@@ -2986,6 +3000,7 @@ var AgentPublicService = /** @class */ (function () {
                                 agentName: agentName,
                                 swarmName: "",
                                 storageName: "",
+                                stateName: "",
                             })];
                     case 1: return [2 /*return*/, _a.sent()];
                 }
@@ -3018,6 +3033,7 @@ var AgentPublicService = /** @class */ (function () {
                                 agentName: agentName,
                                 swarmName: "",
                                 storageName: "",
+                                stateName: "",
                             })];
                     case 1: return [2 /*return*/, _a.sent()];
                 }
@@ -3050,6 +3066,7 @@ var AgentPublicService = /** @class */ (function () {
                                 agentName: agentName,
                                 swarmName: "",
                                 storageName: "",
+                                stateName: "",
                             })];
                     case 1: return [2 /*return*/, _a.sent()];
                 }
@@ -3096,6 +3113,7 @@ var HistoryPublicService = /** @class */ (function () {
                                 agentName: agentName,
                                 swarmName: "",
                                 storageName: "",
+                                stateName: "",
                             })];
                     case 1: return [2 /*return*/, _a.sent()];
                 }
@@ -3130,6 +3148,7 @@ var HistoryPublicService = /** @class */ (function () {
                                 agentName: agentName,
                                 swarmName: "",
                                 storageName: "",
+                                stateName: "",
                             })];
                     case 1: return [2 /*return*/, _a.sent()];
                 }
@@ -3162,6 +3181,7 @@ var HistoryPublicService = /** @class */ (function () {
                                 agentName: agentName,
                                 swarmName: "",
                                 storageName: "",
+                                stateName: "",
                             })];
                     case 1: return [2 /*return*/, _a.sent()];
                 }
@@ -3194,6 +3214,7 @@ var HistoryPublicService = /** @class */ (function () {
                                 agentName: agentName,
                                 swarmName: "",
                                 storageName: "",
+                                stateName: "",
                             })];
                     case 1: return [2 /*return*/, _a.sent()];
                 }
@@ -3240,6 +3261,7 @@ var SessionPublicService = /** @class */ (function () {
                                 swarmName: swarmName,
                                 agentName: "",
                                 storageName: "",
+                                stateName: "",
                             })];
                     case 1: return [2 /*return*/, _a.sent()];
                 }
@@ -3275,6 +3297,7 @@ var SessionPublicService = /** @class */ (function () {
                                 swarmName: swarmName,
                                 agentName: "",
                                 storageName: "",
+                                stateName: "",
                             })];
                     case 1: return [2 /*return*/, _a.sent()];
                 }
@@ -3309,6 +3332,7 @@ var SessionPublicService = /** @class */ (function () {
                                     swarmName: swarmName,
                                     agentName: "",
                                     storageName: "",
+                                    stateName: "",
                                 })];
                             case 1: return [2 /*return*/, _a.sent()];
                         }
@@ -3322,6 +3346,7 @@ var SessionPublicService = /** @class */ (function () {
                         swarmName: swarmName,
                         agentName: "",
                         storageName: "",
+                        stateName: "",
                     });
                 };
             }, {
@@ -3329,6 +3354,7 @@ var SessionPublicService = /** @class */ (function () {
                 swarmName: swarmName,
                 agentName: "",
                 storageName: "",
+                stateName: "",
             });
         };
         /**
@@ -3362,6 +3388,7 @@ var SessionPublicService = /** @class */ (function () {
                                 swarmName: swarmName,
                                 agentName: "",
                                 storageName: "",
+                                stateName: "",
                             })];
                     case 1: return [2 /*return*/, _a.sent()];
                 }
@@ -3396,6 +3423,7 @@ var SessionPublicService = /** @class */ (function () {
                                 swarmName: swarmName,
                                 agentName: "",
                                 storageName: "",
+                                stateName: "",
                             })];
                     case 1: return [2 /*return*/, _a.sent()];
                 }
@@ -3430,6 +3458,7 @@ var SessionPublicService = /** @class */ (function () {
                                 swarmName: swarmName,
                                 agentName: "",
                                 storageName: "",
+                                stateName: "",
                             })];
                     case 1: return [2 /*return*/, _a.sent()];
                 }
@@ -3462,6 +3491,7 @@ var SessionPublicService = /** @class */ (function () {
                                 swarmName: swarmName,
                                 agentName: "",
                                 storageName: "",
+                                stateName: "",
                             })];
                     case 1: return [2 /*return*/, _a.sent()];
                 }
@@ -3494,6 +3524,7 @@ var SessionPublicService = /** @class */ (function () {
                                 swarmName: swarmName,
                                 agentName: "",
                                 storageName: "",
+                                stateName: "",
                             })];
                     case 1: return [2 /*return*/, _a.sent()];
                 }
@@ -3538,6 +3569,7 @@ var SwarmPublicService = /** @class */ (function () {
                                 swarmName: swarmName,
                                 agentName: "",
                                 storageName: "",
+                                stateName: "",
                             })];
                     case 1: return [2 /*return*/, _a.sent()];
                 }
@@ -3570,6 +3602,7 @@ var SwarmPublicService = /** @class */ (function () {
                                 swarmName: swarmName,
                                 agentName: "",
                                 storageName: "",
+                                stateName: "",
                             })];
                     case 1: return [2 /*return*/, _a.sent()];
                 }
@@ -3602,6 +3635,7 @@ var SwarmPublicService = /** @class */ (function () {
                                 swarmName: swarmName,
                                 agentName: "",
                                 storageName: "",
+                                stateName: "",
                             })];
                     case 1: return [2 /*return*/, _a.sent()];
                 }
@@ -3638,6 +3672,7 @@ var SwarmPublicService = /** @class */ (function () {
                                 swarmName: swarmName,
                                 agentName: "",
                                 storageName: "",
+                                stateName: "",
                             })];
                     case 1: return [2 /*return*/, _a.sent()];
                 }
@@ -3672,6 +3707,7 @@ var SwarmPublicService = /** @class */ (function () {
                                 swarmName: swarmName,
                                 agentName: "",
                                 storageName: "",
+                                stateName: "",
                             })];
                     case 1: return [2 /*return*/, _a.sent()];
                 }
@@ -3704,6 +3740,7 @@ var SwarmPublicService = /** @class */ (function () {
                                 swarmName: swarmName,
                                 agentName: "",
                                 storageName: "",
+                                stateName: "",
                             })];
                     case 1: return [2 /*return*/, _a.sent()];
                 }
@@ -3736,6 +3773,18 @@ var AgentValidationService = /** @class */ (function () {
             }
             return _this._agentMap.get(agentName).storages;
         };
+        /**
+         * Retrieves the states used by the agent
+         * @param {agentName} agentName - The name of the swarm.
+         * @returns {string[]} The list of state names.
+         * @throws Will throw an error if the swarm is not found.
+         */
+        this.getStateList = function (agentName) {
+            if (!_this._agentMap.has(agentName)) {
+                throw new Error("agent-swarm agent ".concat(agentName, " not exist (getStateList)"));
+            }
+            return _this._agentMap.get(agentName).states;
+        };
         /**
          * Adds a new agent to the validation service.
          * @param {AgentName} agentName - The name of the agent.
@@ -3769,6 +3818,23 @@ var AgentValidationService = /** @class */ (function () {
             var _a = _this._agentMap.get(agentName).storages, storages = _a === void 0 ? [] : _a;
             return storages.includes(storageName);
         });
+        /**
+         * Check if agent got registered state
+         */
+        this.hasState = functoolsKit.memoize(function (_a) {
+            var _b = __read(_a, 2), agentName = _b[0], stateName = _b[1];
+            return "".concat(agentName, "-").concat(stateName);
+        }, function (agentName, stateName) {
+            _this.loggerService.log("agentValidationService hasState", {
+                agentName: agentName,
+                stateName: stateName,
+            });
+            if (!_this._agentMap.has(agentName)) {
+                throw new Error("agent-swarm agent ".concat(agentName, " not exist (hasState)"));
+            }
+            var _a = _this._agentMap.get(agentName).states, states = _a === void 0 ? [] : _a;
+            return states.includes(stateName);
+        });
         /**
          * Validates an agent by its name and source.
          * @param {AgentName} agentName - The name of the agent.
@@ -3862,6 +3928,7 @@ var SessionValidationService = /** @class */ (function () {
         this._storageSwarmMap = new Map();
         this._historySwarmMap = new Map();
         this._agentSwarmMap = new Map();
+        this._stateSwarmMap = new Map();
         this._sessionSwarmMap = new Map();
         this._sessionModeMap = new Map();
         /**
@@ -3941,6 +4008,26 @@ var SessionValidationService = /** @class */ (function () {
                 _this._storageSwarmMap.set(sessionId, [storageName]);
             }
         };
+        /**
+         * Adds a state usage to a session.
+         * @param {SessionId} sessionId - The ID of the session.
+         * @param {StateName} stateName - The name of the state.
+         */
+        this.addStateUsage = function (sessionId, stateName) {
+            _this.loggerService.log("sessionValidationService addStateUsage", {
+                sessionId: sessionId,
+                stateName: stateName,
+            });
+            if (_this._stateSwarmMap.has(sessionId)) {
+                var states = _this._stateSwarmMap.get(sessionId);
+                if (!states.includes(stateName)) {
+                    states.push(stateName);
+                }
+            }
+            else {
+                _this._stateSwarmMap.set(sessionId, [stateName]);
+            }
+        };
         /**
          * Removes an agent usage from a session.
          * @param {SessionId} sessionId - The ID of the session.
@@ -3992,11 +4079,11 @@ var SessionValidationService = /** @class */ (function () {
             }
         };
         /**
-        * Removes a storage usage from a session.
-        * @param {SessionId} sessionId - The ID of the session.
-        * @param {StorageName} storageName - The name of the storage.
-        * @throws Will throw an error if no storages are found for the session.
-        */
+         * Removes a storage usage from a session.
+         * @param {SessionId} sessionId - The ID of the session.
+         * @param {StorageName} storageName - The name of the storage.
+         * @throws Will throw an error if no storages are found for the session.
+         */
         this.removeStorageUsage = function (sessionId, storageName) {
             _this.loggerService.log("sessionValidationService removeStorageUsage", {
                 sessionId: sessionId,
@@ -4016,6 +4103,31 @@ var SessionValidationService = /** @class */ (function () {
                 throw new Error("No agents found for sessionId=".concat(sessionId));
             }
         };
+        /**
+         * Removes a state usage from a session.
+         * @param {SessionId} sessionId - The ID of the session.
+         * @param {StateName} stateName - The name of the state.
+         * @throws Will throw an error if no states are found for the session.
+         */
+        this.removeStateUsage = function (sessionId, stateName) {
+            _this.loggerService.log("sessionValidationService removeStateUsage", {
+                sessionId: sessionId,
+                stateName: stateName,
+            });
+            if (_this._stateSwarmMap.has(sessionId)) {
+                var agents = _this._stateSwarmMap.get(sessionId);
+                var agentIndex = agents.indexOf(stateName);
+                if (agentIndex !== -1) {
+                    agents.splice(agentIndex, 1);
+                }
+                if (agents.length === 0) {
+                    _this._stateSwarmMap.delete(sessionId);
+                }
+            }
+            else {
+                throw new Error("No agents found for sessionId=".concat(sessionId));
+            }
+        };
         /**
          * Gets the mode of a session.
          * @param {SessionId} clientId - The ID of the client.
@@ -4304,6 +4416,7 @@ var ClientStorage = /** @class */ (function () {
      */
     function ClientStorage(params) {
         var _this = this;
+        var _a;
         this.params = params;
         this._itemMap = new Map();
         /**
@@ -4529,9 +4642,26 @@ var ClientStorage = /** @class */ (function () {
                 return [2 /*return*/, result];
             });
         }); };
+        /**
+         * Disposes of the state.
+         * @returns {Promise<void>}
+         */
+        this.dispose = function () { return __awaiter(_this, void 0, void 0, function () {
+            var _a;
+            return __generator(this, function (_b) {
+                this.params.logger.debug("ClientStorage storageName=".concat(this.params.storageName, " clientId=").concat(this.params.clientId, " shared=").concat(this.params.shared, " dispose"));
+                if ((_a = this.params.callbacks) === null || _a === void 0 ? void 0 : _a.onDispose) {
+                    this.params.callbacks.onDispose(this.params.clientId, this.params.storageName);
+                }
+                return [2 /*return*/];
+            });
+        }); };
         this.params.logger.debug("ClientStorage storageName=".concat(this.params.storageName, " clientId=").concat(this.params.clientId, " shared=").concat(this.params.shared, " CTOR"), {
             params: params,
         });
+        if ((_a = this.params.callbacks) === null || _a === void 0 ? void 0 : _a.onInit) {
+            this.params.callbacks.onInit(this.params.clientId, this.params.storageName);
+        }
     }
     return ClientStorage;
 }());
@@ -4726,18 +4856,31 @@ var StorageConnectionService = /** @class */ (function () {
          * @returns {Promise<void>}
          */
         this.dispose = function () { return __awaiter(_this, void 0, void 0, function () {
-            var key;
+            var key, storage;
             return __generator(this, function (_a) {
-                this.loggerService.log("storageConnectionService dispose", {
-                    context: this.contextService.context,
-                });
-                key = "".concat(this.contextService.context.clientId, "-").concat(this.contextService.context.storageName);
-                if (!this.getStorage.has(key)) {
-                    return [2 /*return*/];
+                switch (_a.label) {
+                    case 0:
+                        this.loggerService.log("storageConnectionService dispose", {
+                            context: this.contextService.context,
+                        });
+                        key = "".concat(this.contextService.context.clientId, "-").concat(this.contextService.context.storageName);
+                        if (!this.getStorage.has(key)) {
+                            return [2 /*return*/];
+                        }
+                        if (!!this.getSharedStorage.has(this.contextService.context.storageName)) return [3 /*break*/, 3];
+                        storage = this.getSharedStorage(this.contextService.context.clientId, this.contextService.context.storageName);
+                        return [4 /*yield*/, storage.waitForInit()];
+                    case 1:
+                        _a.sent();
+                        return [4 /*yield*/, storage.dispose()];
+                    case 2:
+                        _a.sent();
+                        _a.label = 3;
+                    case 3:
+                        this.getStorage.clear(key);
+                        this.sessionValidationService.removeStorageUsage(this.contextService.context.clientId, this.contextService.context.storageName);
+                        return [2 /*return*/];
                 }
-                this.getStorage.clear(key);
-                this.sessionValidationService.removeStorageUsage(this.contextService.context.clientId, this.contextService.context.storageName);
-                return [2 /*return*/];
             });
         }); };
     }
@@ -4782,6 +4925,7 @@ var StoragePublicService = /** @class */ (function () {
                                 storageName: storageName,
                                 agentName: "",
                                 swarmName: "",
+                                stateName: "",
                             })];
                     case 1: return [2 /*return*/, _a.sent()];
                 }
@@ -4814,6 +4958,7 @@ var StoragePublicService = /** @class */ (function () {
                                 storageName: storageName,
                                 agentName: "",
                                 swarmName: "",
+                                stateName: "",
                             })];
                     case 1: return [2 /*return*/, _a.sent()];
                 }
@@ -4846,6 +4991,7 @@ var StoragePublicService = /** @class */ (function () {
                                 storageName: storageName,
                                 agentName: "",
                                 swarmName: "",
+                                stateName: "",
                             })];
                     case 1: return [2 /*return*/, _a.sent()];
                 }
@@ -4878,6 +5024,7 @@ var StoragePublicService = /** @class */ (function () {
                                 storageName: storageName,
                                 agentName: "",
                                 swarmName: "",
+                                stateName: "",
                             })];
                     case 1: return [2 /*return*/, _a.sent()];
                 }
@@ -4909,6 +5056,7 @@ var StoragePublicService = /** @class */ (function () {
                                 storageName: storageName,
                                 agentName: "",
                                 swarmName: "",
+                                stateName: "",
                             })];
                     case 1: return [2 /*return*/, _a.sent()];
                 }
@@ -4939,6 +5087,7 @@ var StoragePublicService = /** @class */ (function () {
                                 storageName: storageName,
                                 agentName: "",
                                 swarmName: "",
+                                stateName: "",
                             })];
                     case 1: return [2 /*return*/, _a.sent()];
                 }
@@ -4971,6 +5120,7 @@ var StoragePublicService = /** @class */ (function () {
                                 storageName: storageName,
                                 agentName: "",
                                 swarmName: "",
+                                stateName: "",
                             })];
                     case 1: return [2 /*return*/, _a.sent()];
                 }
@@ -5077,6 +5227,470 @@ var EmbeddingValidationService = /** @class */ (function () {
     return EmbeddingValidationService;
 }());
 
+/**
+ * Service for managing state schemas.
+ */
+var StateSchemaService = /** @class */ (function () {
+    function StateSchemaService() {
+        var _this = this;
+        this.loggerService = inject(TYPES.loggerService);
+        this.registry = new functoolsKit.ToolRegistry("stateSchemaService");
+        /**
+         * Registers a new state schema.
+         * @param {StateName} key - The key for the schema.
+         * @param {IStateSchema} value - The schema to register.
+         */
+        this.register = function (key, value) {
+            _this.loggerService.log("stateSchemaService register", { key: key });
+            _this.registry = _this.registry.register(key, value);
+        };
+        /**
+         * Retrieves a state schema by key.
+         * @param {StateName} key - The key of the schema to retrieve.
+         * @returns {IStateSchema} The retrieved schema.
+         */
+        this.get = function (key) {
+            _this.loggerService.log("stateSchemaService get", { key: key });
+            return _this.registry.get(key);
+        };
+    }
+    return StateSchemaService;
+}());
+
+/**
+ * Class representing the client state.
+ * @template State - The type of the state data.
+ * @implements {IState<State>}
+ */
+var ClientState = /** @class */ (function () {
+    /**
+     * Creates an instance of ClientState.
+     * @param {IStateParams<State>} params - The state parameters.
+     */
+    function ClientState(params) {
+        var _this = this;
+        var _a;
+        this.params = params;
+        this._state = null;
+        this.dispatch = functoolsKit.queued(function (action, payload) { return __awaiter(_this, void 0, void 0, function () {
+            var _a;
+            return __generator(this, function (_b) {
+                switch (_b.label) {
+                    case 0:
+                        if (action === "read") {
+                            return [2 /*return*/, this._state];
+                        }
+                        if (!(action === "write")) return [3 /*break*/, 2];
+                        console.assert(payload, "agent-swarm ClientState write action undefined payload");
+                        _a = this;
+                        return [4 /*yield*/, payload(this._state)];
+                    case 1: return [2 /*return*/, _a._state = _b.sent()];
+                    case 2: throw new Error("agent-swarm ClientState unknown action");
+                }
+            });
+        }); });
+        /**
+         * Waits for the state to initialize.
+         * @returns {Promise<void>}
+         */
+        this.waitForInit = functoolsKit.singleshot(function () { return __awaiter(_this, void 0, void 0, function () {
+            var _a;
+            var _b;
+            return __generator(this, function (_c) {
+                switch (_c.label) {
+                    case 0:
+                        this.params.logger.debug("ClientState stateName=".concat(this.params.stateName, " clientId=").concat(this.params.clientId, " shared=").concat(this.params.shared, " waitForInit"));
+                        _a = this;
+                        return [4 /*yield*/, this.params.getState(this.params.clientId, this.params.stateName)];
+                    case 1:
+                        _a._state = _c.sent();
+                        this.params.logger.debug("ClientState stateName=".concat(this.params.stateName, " clientId=").concat(this.params.clientId, " shared=").concat(this.params.shared, " waitForInit output"), { initialState: this._state });
+                        if ((_b = this.params.callbacks) === null || _b === void 0 ? void 0 : _b.onLoad) {
+                            this.params.callbacks.onLoad(this._state, this.params.clientId, this.params.stateName);
+                        }
+                        return [2 /*return*/];
+                }
+            });
+        }); });
+        /**
+         * Sets the state using the provided dispatch function.
+         * @param {DispatchFn<State>} dispatchFn - The dispatch function.
+         * @returns {Promise<State>}
+         */
+        this.setState = function (dispatchFn) { return __awaiter(_this, void 0, void 0, function () {
+            var _this = this;
+            var _a;
+            return __generator(this, function (_b) {
+                switch (_b.label) {
+                    case 0:
+                        this.params.logger.debug("ClientState stateName=".concat(this.params.stateName, " clientId=").concat(this.params.clientId, " shared=").concat(this.params.shared, " setState"));
+                        return [4 /*yield*/, this.dispatch("write", function (currentState) { return __awaiter(_this, void 0, void 0, function () {
+                                var _a, _b, middleware, e_1_1;
+                                var e_1, _c;
+                                return __generator(this, function (_d) {
+                                    switch (_d.label) {
+                                        case 0:
+                                            _d.trys.push([0, 5, 6, 7]);
+                                            _a = __values(this.params.middlewares), _b = _a.next();
+                                            _d.label = 1;
+                                        case 1:
+                                            if (!!_b.done) return [3 /*break*/, 4];
+                                            middleware = _b.value;
+                                            return [4 /*yield*/, middleware(currentState, this.params.clientId, this.params.stateName)];
+                                        case 2:
+                                            currentState = _d.sent();
+                                            _d.label = 3;
+                                        case 3:
+                                            _b = _a.next();
+                                            return [3 /*break*/, 1];
+                                        case 4: return [3 /*break*/, 7];
+                                        case 5:
+                                            e_1_1 = _d.sent();
+                                            e_1 = { error: e_1_1 };
+                                            return [3 /*break*/, 7];
+                                        case 6:
+                                            try {
+                                                if (_b && !_b.done && (_c = _a.return)) _c.call(_a);
+                                            }
+                                            finally { if (e_1) throw e_1.error; }
+                                            return [7 /*endfinally*/];
+                                        case 7: return [4 /*yield*/, dispatchFn(currentState)];
+                                        case 8: return [2 /*return*/, _d.sent()];
+                                    }
+                                });
+                            }); })];
+                    case 1:
+                        _b.sent();
+                        this.params.logger.debug("ClientState stateName=".concat(this.params.stateName, " clientId=").concat(this.params.clientId, " shared=").concat(this.params.shared, " setState output"), { pendingState: this._state });
+                        this.params.setState &&
+                            this.params.setState(this._state, this.params.clientId, this.params.stateName);
+                        if ((_a = this.params.callbacks) === null || _a === void 0 ? void 0 : _a.onWrite) {
+                            this.params.callbacks.onWrite(this._state, this.params.clientId, this.params.stateName);
+                        }
+                        return [2 /*return*/, this._state];
+                }
+            });
+        }); };
+        /**
+         * Gets the current state.
+         * @returns {Promise<State>}
+         */
+        this.getState = function () { return __awaiter(_this, void 0, void 0, function () {
+            var _a;
+            return __generator(this, function (_b) {
+                switch (_b.label) {
+                    case 0:
+                        this.params.logger.debug("ClientState stateName=".concat(this.params.stateName, " clientId=").concat(this.params.clientId, " shared=").concat(this.params.shared, " getState"));
+                        return [4 /*yield*/, this.dispatch("read")];
+                    case 1:
+                        _b.sent();
+                        if ((_a = this.params.callbacks) === null || _a === void 0 ? void 0 : _a.onRead) {
+                            this.params.callbacks.onRead(this._state, this.params.clientId, this.params.stateName);
+                        }
+                        return [2 /*return*/, this._state];
+                }
+            });
+        }); };
+        /**
+         * Disposes of the state.
+         * @returns {Promise<void>}
+         */
+        this.dispose = function () { return __awaiter(_this, void 0, void 0, function () {
+            var _a;
+            return __generator(this, function (_b) {
+                this.params.logger.debug("ClientState stateName=".concat(this.params.stateName, " clientId=").concat(this.params.clientId, " shared=").concat(this.params.shared, " dispose"));
+                if ((_a = this.params.callbacks) === null || _a === void 0 ? void 0 : _a.onDispose) {
+                    this.params.callbacks.onDispose(this.params.clientId, this.params.stateName);
+                }
+                return [2 /*return*/];
+            });
+        }); };
+        this.params.logger.debug("ClientState stateName=".concat(this.params.stateName, " clientId=").concat(this.params.clientId, " shared=").concat(this.params.shared, " CTOR"), {
+            params: params,
+        });
+        if ((_a = this.params.callbacks) === null || _a === void 0 ? void 0 : _a.onInit) {
+            this.params.callbacks.onInit(this.params.clientId, this.params.stateName);
+        }
+    }
+    return ClientState;
+}());
+
+/**
+ * Service for managing state connections.
+ * @template T - The type of state data.
+ * @implements {IState<T>}
+ */
+var StateConnectionService = /** @class */ (function () {
+    function StateConnectionService() {
+        var _this = this;
+        this.loggerService = inject(TYPES.loggerService);
+        this.contextService = inject(TYPES.contextService);
+        this.stateSchemaService = inject(TYPES.stateSchemaService);
+        this.sessionValidationService = inject(TYPES.sessionValidationService);
+        /**
+         * Memoized function to get a shared state reference.
+         * @param {string} clientId - The client ID.
+         * @param {StateName} stateName - The state name.
+         * @returns {ClientState} The client state.
+         */
+        this.getSharedStateRef = functoolsKit.memoize(function (_a) {
+            var _b = __read(_a, 2), stateName = _b[1];
+            return "".concat(stateName);
+        }, function (clientId, stateName) {
+            _this.sessionValidationService.addStateUsage(clientId, stateName);
+            var _a = _this.stateSchemaService.get(stateName), getState = _a.getState, setState = _a.setState, _b = _a.middlewares, middlewares = _b === void 0 ? [] : _b, shared = _a.shared, callbacks = _a.callbacks;
+            if (!shared) {
+                throw new Error("agent-swarm state not shared stateName=".concat(stateName));
+            }
+            return new ClientState({
+                clientId: clientId,
+                stateName: stateName,
+                logger: _this.loggerService,
+                setState: setState
+                    ? functoolsKit.queued(function () {
+                        var args = [];
+                        for (var _i = 0; _i < arguments.length; _i++) {
+                            args[_i] = arguments[_i];
+                        }
+                        return __awaiter(_this, void 0, void 0, function () { return __generator(this, function (_a) {
+                            switch (_a.label) {
+                                case 0: return [4 /*yield*/, setState.apply(void 0, __spreadArray([], __read(args), false))];
+                                case 1: return [2 /*return*/, _a.sent()];
+                            }
+                        }); });
+                    })
+                    : setState,
+                getState: getState,
+                middlewares: middlewares,
+                callbacks: callbacks,
+            });
+        });
+        /**
+         * Memoized function to get a state reference.
+         * @param {string} clientId - The client ID.
+         * @param {StateName} stateName - The state name.
+         * @returns {ClientState} The client state.
+         */
+        this.getStateRef = functoolsKit.memoize(function (_a) {
+            var _b = __read(_a, 2), clientId = _b[0], stateName = _b[1];
+            return "".concat(clientId, "-").concat(stateName);
+        }, function (clientId, stateName) {
+            _this.sessionValidationService.addStateUsage(clientId, stateName);
+            var _a = _this.stateSchemaService.get(stateName), getState = _a.getState, setState = _a.setState, _b = _a.middlewares, middlewares = _b === void 0 ? [] : _b, callbacks = _a.callbacks, _c = _a.shared, shared = _c === void 0 ? false : _c;
+            if (shared) {
+                return _this.getSharedStateRef(clientId, stateName);
+            }
+            return new ClientState({
+                clientId: clientId,
+                stateName: stateName,
+                logger: _this.loggerService,
+                setState: setState
+                    ? functoolsKit.queued(function () {
+                        var args = [];
+                        for (var _i = 0; _i < arguments.length; _i++) {
+                            args[_i] = arguments[_i];
+                        }
+                        return __awaiter(_this, void 0, void 0, function () { return __generator(this, function (_a) {
+                            switch (_a.label) {
+                                case 0: return [4 /*yield*/, setState.apply(void 0, __spreadArray([], __read(args), false))];
+                                case 1: return [2 /*return*/, _a.sent()];
+                            }
+                        }); });
+                    })
+                    : setState,
+                getState: getState,
+                middlewares: middlewares,
+                callbacks: callbacks,
+            });
+        });
+        /**
+         * Sets the state.
+         * @param {function(T): Promise<T>} dispatchFn - The function to dispatch the new state.
+         * @returns {Promise<T>} The new state.
+         */
+        this.setState = function (dispatchFn) { return __awaiter(_this, void 0, void 0, function () {
+            var state;
+            return __generator(this, function (_a) {
+                switch (_a.label) {
+                    case 0:
+                        this.loggerService.log("stateConnectionService setState", {
+                            context: this.contextService.context,
+                        });
+                        state = this.getStateRef(this.contextService.context.clientId, this.contextService.context.stateName);
+                        return [4 /*yield*/, state.waitForInit()];
+                    case 1:
+                        _a.sent();
+                        return [4 /*yield*/, state.setState(dispatchFn)];
+                    case 2: return [2 /*return*/, _a.sent()];
+                }
+            });
+        }); };
+        /**
+         * Gets the state.
+         * @returns {Promise<T>} The current state.
+         */
+        this.getState = function () { return __awaiter(_this, void 0, void 0, function () {
+            var state;
+            return __generator(this, function (_a) {
+                switch (_a.label) {
+                    case 0:
+                        this.loggerService.log("stateConnectionService getState", {
+                            context: this.contextService.context,
+                        });
+                        state = this.getStateRef(this.contextService.context.clientId, this.contextService.context.stateName);
+                        return [4 /*yield*/, state.waitForInit()];
+                    case 1:
+                        _a.sent();
+                        return [4 /*yield*/, state.getState()];
+                    case 2: return [2 /*return*/, _a.sent()];
+                }
+            });
+        }); };
+        /**
+         * Disposes the state connection.
+         * @returns {Promise<void>}
+         */
+        this.dispose = function () { return __awaiter(_this, void 0, void 0, function () {
+            var key, state;
+            return __generator(this, function (_a) {
+                switch (_a.label) {
+                    case 0:
+                        this.loggerService.log("stateConnectionService dispose", {
+                            context: this.contextService.context,
+                        });
+                        key = "".concat(this.contextService.context.clientId, "-").concat(this.contextService.context.stateName);
+                        if (!this.getStateRef.has(key)) {
+                            return [2 /*return*/];
+                        }
+                        if (!!this.getSharedStateRef.has(this.contextService.context.stateName)) return [3 /*break*/, 3];
+                        state = this.getStateRef(this.contextService.context.clientId, this.contextService.context.stateName);
+                        return [4 /*yield*/, state.waitForInit()];
+                    case 1:
+                        _a.sent();
+                        return [4 /*yield*/, state.dispose()];
+                    case 2:
+                        _a.sent();
+                        _a.label = 3;
+                    case 3:
+                        this.getStateRef.clear(key);
+                        this.sessionValidationService.removeStateUsage(this.contextService.context.clientId, this.contextService.context.stateName);
+                        return [2 /*return*/];
+                }
+            });
+        }); };
+    }
+    return StateConnectionService;
+}());
+
+var StatePublicService = /** @class */ (function () {
+    function StatePublicService() {
+        var _this = this;
+        this.loggerService = inject(TYPES.loggerService);
+        this.stateConnectionService = inject(TYPES.stateConnectionService);
+        /**
+         * Sets the state using the provided dispatch function.
+         * @param {function(T): Promise<T>} dispatchFn - The function to dispatch the state change.
+         * @param {string} clientId - The client ID.
+         * @param {StateName} stateName - The name of the state.
+         * @returns {Promise<T>} - The updated state.
+         */
+        this.setState = function (dispatchFn, clientId, stateName) { return __awaiter(_this, void 0, void 0, function () {
+            var _this = this;
+            return __generator(this, function (_a) {
+                switch (_a.label) {
+                    case 0:
+                        this.loggerService.log("statePublicService setState", {
+                            clientId: clientId,
+                            stateName: stateName,
+                        });
+                        return [4 /*yield*/, ContextService.runInContext(function () { return __awaiter(_this, void 0, void 0, function () {
+                                return __generator(this, function (_a) {
+                                    switch (_a.label) {
+                                        case 0: return [4 /*yield*/, this.stateConnectionService.setState(dispatchFn)];
+                                        case 1: return [2 /*return*/, _a.sent()];
+                                    }
+                                });
+                            }); }, {
+                                clientId: clientId,
+                                stateName: stateName,
+                                agentName: "",
+                                swarmName: "",
+                                storageName: "",
+                            })];
+                    case 1: return [2 /*return*/, _a.sent()];
+                }
+            });
+        }); };
+        /**
+         * Gets the current state.
+         * @param {string} clientId - The client ID.
+         * @param {StateName} stateName - The name of the state.
+         * @returns {Promise<T>} - The current state.
+         */
+        this.getState = function (clientId, stateName) { return __awaiter(_this, void 0, void 0, function () {
+            var _this = this;
+            return __generator(this, function (_a) {
+                switch (_a.label) {
+                    case 0:
+                        this.loggerService.log("statePublicService getState", {
+                            clientId: clientId,
+                            stateName: stateName,
+                        });
+                        return [4 /*yield*/, ContextService.runInContext(function () { return __awaiter(_this, void 0, void 0, function () {
+                                return __generator(this, function (_a) {
+                                    switch (_a.label) {
+                                        case 0: return [4 /*yield*/, this.stateConnectionService.getState()];
+                                        case 1: return [2 /*return*/, _a.sent()];
+                                    }
+                                });
+                            }); }, {
+                                clientId: clientId,
+                                stateName: stateName,
+                                agentName: "",
+                                swarmName: "",
+                                storageName: "",
+                            })];
+                    case 1: return [2 /*return*/, _a.sent()];
+                }
+            });
+        }); };
+        /**
+         * Disposes the state.
+         * @param {string} clientId - The client ID.
+         * @param {StateName} stateName - The name of the state.
+         * @returns {Promise<void>} - A promise that resolves when the state is disposed.
+         */
+        this.dispose = function (clientId, stateName) { return __awaiter(_this, void 0, void 0, function () {
+            var _this = this;
+            return __generator(this, function (_a) {
+                switch (_a.label) {
+                    case 0:
+                        this.loggerService.log("statePublicService dispose", {
+                            clientId: clientId,
+                            stateName: stateName,
+                        });
+                        return [4 /*yield*/, ContextService.runInContext(function () { return __awaiter(_this, void 0, void 0, function () {
+                                return __generator(this, function (_a) {
+                                    switch (_a.label) {
+                                        case 0: return [4 /*yield*/, this.stateConnectionService.dispose()];
+                                        case 1: return [2 /*return*/, _a.sent()];
+                                    }
+                                });
+                            }); }, {
+                                clientId: clientId,
+                                stateName: stateName,
+                                agentName: "",
+                                swarmName: "",
+                                storageName: "",
+                            })];
+                    case 1: return [2 /*return*/, _a.sent()];
+                }
+            });
+        }); };
+    }
+    return StatePublicService;
+}());
+
 {
     provide(TYPES.loggerService, function () { return new LoggerService(); });
     provide(TYPES.contextService, function () { return new ContextService(); });
@@ -5087,6 +5701,7 @@ var EmbeddingValidationService = /** @class */ (function () {
     provide(TYPES.swarmConnectionService, function () { return new SwarmConnectionService(); });
     provide(TYPES.sessionConnectionService, function () { return new SessionConnectionService(); });
     provide(TYPES.storageConnectionService, function () { return new StorageConnectionService(); });
+    provide(TYPES.stateConnectionService, function () { return new StateConnectionService(); });
 }
 {
     provide(TYPES.agentSchemaService, function () { return new AgentSchemaService(); });
@@ -5095,6 +5710,7 @@ var EmbeddingValidationService = /** @class */ (function () {
     provide(TYPES.completionSchemaService, function () { return new CompletionSchemaService(); });
     provide(TYPES.embeddingSchemaService, function () { return new EmbeddingSchemaService(); });
     provide(TYPES.storageSchemaService, function () { return new StorageSchemaService(); });
+    provide(TYPES.stateSchemaService, function () { return new StateSchemaService(); });
 }
 {
     provide(TYPES.agentPublicService, function () { return new AgentPublicService(); });
@@ -5102,6 +5718,7 @@ var EmbeddingValidationService = /** @class */ (function () {
     provide(TYPES.sessionPublicService, function () { return new SessionPublicService(); });
     provide(TYPES.swarmPublicService, function () { return new SwarmPublicService(); });
     provide(TYPES.storagePublicService, function () { return new StoragePublicService(); });
+    provide(TYPES.statePublicService, function () { return new StatePublicService(); });
 }
 {
     provide(TYPES.agentPublicService, function () { return new AgentPublicService(); });
@@ -5129,6 +5746,7 @@ var connectionServices = {
     swarmConnectionService: inject(TYPES.swarmConnectionService),
     sessionConnectionService: inject(TYPES.sessionConnectionService),
     storageConnectionService: inject(TYPES.storageConnectionService),
+    stateConnectionService: inject(TYPES.stateConnectionService),
 };
 var schemaServices = {
     agentSchemaService: inject(TYPES.agentSchemaService),
@@ -5137,6 +5755,7 @@ var schemaServices = {
     completionSchemaService: inject(TYPES.completionSchemaService),
     embeddingSchemaService: inject(TYPES.embeddingSchemaService),
     storageSchemaService: inject(TYPES.storageSchemaService),
+    stateSchemaService: inject(TYPES.stateSchemaService),
 };
 var publicServices = {
     agentPublicService: inject(TYPES.agentPublicService),
@@ -5144,6 +5763,7 @@ var publicServices = {
     sessionPublicService: inject(TYPES.sessionPublicService),
     swarmPublicService: inject(TYPES.swarmPublicService),
     storagePublicService: inject(TYPES.storagePublicService),
+    statePublicService: inject(TYPES.statePublicService),
 };
 var validationServices = {
     agentValidationService: inject(TYPES.agentValidationService),
@@ -5219,6 +5839,25 @@ var addTool = function (toolSchema) {
     return toolSchema.toolName;
 };
 
+/**
+ * Adds a new state to the state registry. The swarm takes only those states which was registered
+ *
+ * @param {IStateSchema} stateSchema - The schema of the state to be added.
+ * @returns {string} The name of the added state.
+ */
+var addState = function (stateSchema) {
+    swarm.loggerService.log("function addState", {
+        stateSchema: stateSchema,
+    });
+    swarm.stateSchemaService.register(stateSchema.stateName, stateSchema);
+    if (stateSchema.shared) {
+        swarm.stateConnectionService
+            .getSharedStateRef("shared", stateSchema.stateName)
+            .waitForInit();
+    }
+    return stateSchema.stateName;
+};
+
 /**
  * Adds a new embedding to the embedding registry. The swarm takes only those embeddings which was registered
  *
@@ -5586,6 +6225,7 @@ var disposeConnection = function (clientId, swarmName) { return __awaiter(void 0
                             .flatMap(function (agentName) {
                             return swarm.agentValidationService.getStorageList(agentName);
                         })
+                            .filter(function (storageName) { return !!storageName; })
                             .map(function (storageName) { return __awaiter(void 0, void 0, void 0, function () {
                             return __generator(this, function (_a) {
                                 switch (_a.label) {
@@ -5599,8 +6239,28 @@ var disposeConnection = function (clientId, swarmName) { return __awaiter(void 0
                     ])];
             case 4:
                 _a.sent();
-                return [4 /*yield*/, HistoryAdapter.dispose(clientId, null)];
+                return [4 /*yield*/, Promise.all([
+                        swarm.swarmValidationService
+                            .getAgentList(swarmName)
+                            .flatMap(function (agentName) {
+                            return swarm.agentValidationService.getStateList(agentName);
+                        })
+                            .filter(function (stateName) { return !!stateName; })
+                            .map(function (stateName) { return __awaiter(void 0, void 0, void 0, function () {
+                            return __generator(this, function (_a) {
+                                switch (_a.label) {
+                                    case 0: return [4 /*yield*/, swarm.statePublicService.dispose(clientId, stateName)];
+                                    case 1:
+                                        _a.sent();
+                                        return [2 /*return*/];
+                                }
+                            });
+                        }); }),
+                    ])];
             case 5:
+                _a.sent();
+                return [4 /*yield*/, HistoryAdapter.dispose(clientId, null)];
+            case 6:
                 _a.sent();
                 return [2 /*return*/];
         }
@@ -6614,14 +7274,91 @@ var StorageUtils = /** @class */ (function () {
 }());
 var Storage = new StorageUtils();
 
+/**
+ * Utility class for managing state in the agent swarm.
+ * @implements {TState}
+ */
+var StateUtils = /** @class */ (function () {
+    function StateUtils() {
+        var _this = this;
+        /**
+         * Retrieves the state for a given client and state name.
+         * @template T
+         * @param {Object} payload - The payload containing client and state information.
+         * @param {string} payload.clientId - The client ID.
+         * @param {AgentName} payload.agentName - The agent name.
+         * @param {StateName} payload.stateName - The state name.
+         * @returns {Promise<T>} The state data.
+         * @throws Will throw an error if the state is not registered in the agent.
+         */
+        this.getState = function (payload) { return __awaiter(_this, void 0, void 0, function () {
+            return __generator(this, function (_a) {
+                switch (_a.label) {
+                    case 0:
+                        swarm.loggerService.log("StateUtils getState", {
+                            clientId: payload.clientId,
+                            stateName: payload.stateName,
+                        });
+                        if (!swarm.agentValidationService.hasState(payload.agentName, payload.stateName)) {
+                            throw new Error("agent-swarm StateUtils ".concat(payload.stateName, " not registered in ").concat(payload.agentName, " (getState)"));
+                        }
+                        return [4 /*yield*/, swarm.statePublicService.getState(payload.clientId, payload.stateName)];
+                    case 1: return [2 /*return*/, _a.sent()];
+                }
+            });
+        }); };
+        /**
+         * Sets the state for a given client and state name.
+         * @template T
+         * @param {T | ((prevState: T) => Promise<T>)} dispatchFn - The new state or a function that returns the new state.
+         * @param {Object} payload - The payload containing client and state information.
+         * @param {string} payload.clientId - The client ID.
+         * @param {AgentName} payload.agentName - The agent name.
+         * @param {StateName} payload.stateName - The state name.
+         * @returns {Promise<void>}
+         * @throws Will throw an error if the state is not registered in the agent.
+         */
+        this.setState = function (dispatchFn, payload) { return __awaiter(_this, void 0, void 0, function () {
+            var _this = this;
+            return __generator(this, function (_a) {
+                switch (_a.label) {
+                    case 0:
+                        swarm.loggerService.log("StateUtils setState", {
+                            clientId: payload.clientId,
+                            stateName: payload.stateName,
+                        });
+                        if (!swarm.agentValidationService.hasState(payload.agentName, payload.stateName)) {
+                            throw new Error("agent-swarm StateUtils ".concat(payload.stateName, " not registered in ").concat(payload.agentName, " (setState)"));
+                        }
+                        if (!(typeof dispatchFn === "function")) return [3 /*break*/, 2];
+                        return [4 /*yield*/, swarm.statePublicService.setState(dispatchFn, payload.clientId, payload.stateName)];
+                    case 1: return [2 /*return*/, _a.sent()];
+                    case 2: return [4 /*yield*/, swarm.statePublicService.setState(function () { return __awaiter(_this, void 0, void 0, function () { return __generator(this, function (_a) {
+                            return [2 /*return*/, dispatchFn];
+                        }); }); }, payload.clientId, payload.stateName)];
+                    case 3: return [2 /*return*/, _a.sent()];
+                }
+            });
+        }); };
+    }
+    return StateUtils;
+}());
+/**
+ * Instance of StateUtils for managing state.
+ * @type {StateUtils}
+ */
+var State = new StateUtils();
+
 exports.ContextService = ContextService;
 exports.History = History;
+exports.HistoryAdapter = HistoryAdapter;
+exports.HistoryInstance = HistoryInstance;
+exports.State = State;
 exports.Storage = Storage;
-exports._HistoryAdapter = HistoryAdapter;
-exports._HistoryInstance = HistoryInstance;
 exports.addAgent = addAgent;
 exports.addCompletion = addCompletion;
 exports.addEmbedding = addEmbedding;
+exports.addState = addState;
 exports.addStorage = addStorage;
 exports.addSwarm = addSwarm;
 exports.addTool = addTool;