From b527faa325675dd9a87e811cec71b291731943cb Mon Sep 17 00:00:00 2001 From: Ani Date: Sun, 21 Jun 2026 02:06:48 +0000 Subject: [PATCH 1/4] refactor: introduce StateManagement interface and CompositeStateManagement --- AGENTS.md | 10 +- README.md | 2 +- src/composite-state-management.test.ts | 132 +++++++++++++++++ src/composite-state-management.ts | 194 +++++++++++++++++++++++++ src/hooks.ts | 8 +- src/plugin.ts | 4 +- src/state-interface.ts | 76 ++++++++++ src/state.ts | 73 ++-------- 8 files changed, 423 insertions(+), 76 deletions(-) create mode 100644 src/composite-state-management.test.ts create mode 100644 src/composite-state-management.ts create mode 100644 src/state-interface.ts diff --git a/AGENTS.md b/AGENTS.md index 5881a28..91bf277 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -18,7 +18,7 @@ src/ ├── tool.ts # Sequential Thinking core logic ├── tool-metadata.ts # Tool description and prompt content ├── schema.ts # Tool parameter validation schema -├── state.ts # Session state management +├── state.ts # State management interface and composite implementation ├── config.ts # Plugin configuration processing ├── plugin.test.ts # Plugin integration tests ├── state.test.ts # State management unit tests @@ -36,7 +36,7 @@ index.ts ├─ config.ts → resolveConfig() ├─ tool.ts → SequentialThinkingTool class ├─ schema.ts → TOOL_PARAMETER_SCHEMA - ├─ state.ts → SessionStateManager class + ├─ state.ts → StateManagement interface and CompositeStateManagement implementation ├─ tool-metadata.ts → TOOL_DESCRIPTION & PREFER_SEQUENTIAL_THINKING_CONTEXT └─ hooks.ts → createHookHandlers() ├─ state.ts (for session management) @@ -53,7 +53,7 @@ index.ts | `src/schema.ts` | TypeBox schema for the public `sequential_thinking` tool input contract | | `src/tool-metadata.ts` | Tool description and prompt-injection system context for targeted models | | `src/tool.ts` | `SequentialThinkingTool` class — core thought processing with input validation & no mutation | -| `src/state.ts` | `SessionStateManager` class — encapsulated state lifecycle with SDK cleanup integration | +| `src/state.ts` | `StateManagement` interface and `CompositeStateManagement` class — encapsulated state lifecycle with SDK cleanup integration | | `src/hooks.ts` | SDK hook handlers — manages per-session state mapping and lifecycle events | | `src/plugin.ts` | Plugin orchestration — resolves config, registers tool, session extension, and SDK hooks | @@ -61,13 +61,13 @@ index.ts 1. **Registration Phase**: `registerSequentialThinkingPlugin()` registers tool and hooks 2. **Tool Execution**: `SequentialThinkingTool` processes thinking steps -3. **State Management**: `SessionStateManager` tracks thought history and branches +3. **State Management**: `CompositeStateManagement` tracks thought history and branches 4. **Hooks Processing**: `createHookHandlers()` handles session lifecycle events ### Key Classes - `SequentialThinkingTool`: Core thinking logic, handles `thought`, `branch`, `revision` operations -- `SessionStateManager`: Session state management, supports thought history and branching +- `CompositeStateManagement`: Session state management, supports thought history and branching through interface abstraction - `RunState`: Execution state type containing `thoughtHistory` and `branches` ## Code Style & Patterns diff --git a/README.md b/README.md index a0978ea..0453600 100644 --- a/README.md +++ b/README.md @@ -26,7 +26,7 @@ index.ts ├─ config.ts → resolveConfig() ├─ tool.ts → SequentialThinkingTool class ├─ schema.ts → TOOL_PARAMETER_SCHEMA - ├─ state.ts → SessionStateManager class + ├─ state.ts → StateManagement interface and CompositeStateManagement implementation ├─ tool-metadata.ts → TOOL_DESCRIPTION & PREFER_SEQUENTIAL_THINKING_CONTEXT └─ hooks.ts → createHookHandlers() ├─ state.ts (for session management) diff --git a/src/composite-state-management.test.ts b/src/composite-state-management.test.ts new file mode 100644 index 0000000..392c2f8 --- /dev/null +++ b/src/composite-state-management.test.ts @@ -0,0 +1,132 @@ +import { describe, it, expect, beforeEach } from "vitest"; +import { CompositeStateManagement } from "./composite-state-management.js"; +import { ThoughtData } from "./tool.js"; + +describe("CompositeStateManagement", () => { + let manager: CompositeStateManagement; + + beforeEach(() => { + manager = new CompositeStateManagement(); + }); + + it("should create a new instance", () => { + expect(manager).toBeInstanceOf(CompositeStateManagement); + }); + + it("should register tool calls and associate them with session keys", () => { + manager.registerToolCall("session1", "tool1"); + + // Check that the tool call is registered + const state = manager.getStateByToolCallId("tool1"); + expect(state).toBeDefined(); + expect(state?.thoughtHistory).toEqual([]); + expect(state?.branches).toEqual({}); + }); + + it("should get or create state for a session", () => { + const state = manager.getOrCreateState("session1"); + expect(state).toBeDefined(); + expect(state.thoughtHistory).toEqual([]); + expect(state.branches).toEqual({}); + + // Add some data to the state + const thought: ThoughtData = { + thought: "Test thought", + thoughtNumber: 1, + totalThoughts: 1, + nextThoughtNeeded: false, + }; + + manager.addThought("session1", thought); + + const updatedState = manager.getOrCreateState("session1"); + expect(updatedState.thoughtHistory).toHaveLength(1); + expect(updatedState.thoughtHistory[0].thought).toBe("Test thought"); + }); + + it("should handle branch operations", () => { + const thoughts: ThoughtData[] = [{ + thought: "Branch thought", + thoughtNumber: 1, + totalThoughts: 1, + nextThoughtNeeded: false, + }]; + + manager.addBranch("session1", "branch1", thoughts); + + const state = manager.getOrCreateState("session1"); + expect(state.branches["branch1"]).toBeDefined(); + expect(state.branches["branch1"]).toHaveLength(1); + expect(state.branches["branch1"][0].thought).toBe("Branch thought"); + + const branchIds = manager.getBranchIds("session1"); + expect(branchIds).toContain("branch1"); + }); + + it("should purge session state correctly", () => { + // Add some state + manager.registerToolCall("session1", "tool1"); + const thought: ThoughtData = { + thought: "Test thought", + thoughtNumber: 1, + totalThoughts: 1, + nextThoughtNeeded: false, + }; + manager.addThought("session1", thought); + + // Verify state exists + expect(manager.hasState("session1")).toBe(true); + expect(manager.getOrCreateState("session1").thoughtHistory).toHaveLength(1); + + // Purge the session + manager.purgeSessionState("session1"); + + // Verify state is cleared + expect(manager.hasState("session1")).toBe(false); + const newState = manager.getOrCreateState("session1"); + expect(newState.thoughtHistory).toEqual([]); + expect(newState.branches).toEqual({}); + }); + + it("should track state count", () => { + expect(manager.stateCount).toBe(0); + + manager.getOrCreateState("session1"); + expect(manager.stateCount).toBe(1); + + manager.getOrCreateState("session2"); + expect(manager.stateCount).toBe(2); + + manager.purgeSessionState("session1"); + expect(manager.stateCount).toBe(1); + }); + + it("should handle tool call mapping removal", () => { + manager.registerToolCall("session1", "tool1"); + expect(manager.getStateByToolCallId("tool1")).toBeDefined(); + + manager.removeToolCallMapping("tool1"); + expect(manager.getStateByToolCallId("tool1")).toBeUndefined(); + }); + + it("should provide cleanup callback", () => { + const cleanupFn = manager.getCleanupCallback(); + expect(typeof cleanupFn).toBe("function"); + + // Add some state + const thought: ThoughtData = { + thought: "Test thought", + thoughtNumber: 1, + totalThoughts: 1, + nextThoughtNeeded: false, + }; + manager.addThought("session1", thought); + expect(manager.stateCount).toBe(1); + + // Execute cleanup + cleanupFn("reset"); + + // State should be cleared + expect(manager.stateCount).toBe(0); + }); +}); \ No newline at end of file diff --git a/src/composite-state-management.ts b/src/composite-state-management.ts new file mode 100644 index 0000000..011e66b --- /dev/null +++ b/src/composite-state-management.ts @@ -0,0 +1,194 @@ +import { RunState, ThoughtData } from './tool.js'; +import { StateManagement, StateOperations } from './state-interface.js'; + +/** + * Mapping service for associating tool calls with session keys + */ +class SessionMappingService { + private sessionKeyByToolCallId: Map; + + constructor() { + this.sessionKeyByToolCallId = new Map(); + } + + registerToolCall(sessionKey: string, toolCallId: string): void { + this.sessionKeyByToolCallId.set(toolCallId, sessionKey); + } + + getSessionKey(toolCallId: string): string | undefined { + return this.sessionKeyByToolCallId.get(toolCallId); + } + + removeToolCallMapping(toolCallId: string): void { + this.sessionKeyByToolCallId.delete(toolCallId); + } + + clear(): void { + this.sessionKeyByToolCallId.clear(); + } +} + +/** + * Storage service for maintaining state data + */ +class StateStorageService { + private stateBySessionKey: Map; + + constructor() { + this.stateBySessionKey = new Map(); + } + + getOrCreateState(sessionKey: string): RunState { + if (!this.stateBySessionKey.has(sessionKey)) { + this.stateBySessionKey.set(sessionKey, { + thoughtHistory: [], + branches: {}, + }); + } + return this.stateBySessionKey.get(sessionKey)!; + } + + getState(sessionKey: string): RunState | undefined { + return this.stateBySessionKey.get(sessionKey); + } + + setState(sessionKey: string, state: RunState): void { + this.stateBySessionKey.set(sessionKey, state); + } + + hasState(sessionKey: string): boolean { + return this.stateBySessionKey.has(sessionKey); + } + + purgeSessionState(sessionKey: string): void { + this.stateBySessionKey.delete(sessionKey); + } + + getAllSessionKeys(): string[] { + return Array.from(this.stateBySessionKey.keys()); + } + + clear(): void { + this.stateBySessionKey.clear(); + } + + get stateCount(): number { + return this.stateBySessionKey.size; + } +} + +/** + * Lifecycle management service for handling cleanup operations + */ +class LifecycleManager { + private storageService: StateStorageService; + + constructor(storageService: StateStorageService) { + this.storageService = storageService; + } + + getCleanupCallback(): (action: 'disable' | 'reset' | 'delete' | 'restart') => void { + return (action: 'disable' | 'reset' | 'delete' | 'restart') => { + for (const key of this.storageService.getAllSessionKeys()) { + this.storageService.purgeSessionState(key); + } + }; + } +} + +/** + * Composite implementation of StateManagement and StateOperations + */ +export class CompositeStateManagement implements StateManagement, StateOperations { + private mappingService: SessionMappingService; + private storageService: StateStorageService; + private lifecycleManager: LifecycleManager; + + constructor() { + this.mappingService = new SessionMappingService(); + this.storageService = new StateStorageService(); + this.lifecycleManager = new LifecycleManager(this.storageService); + } + + registerToolCall(sessionKey: string, toolCallId: string): void { + this.mappingService.registerToolCall(sessionKey, toolCallId); + // Ensure state exists for the session + this.storageService.getOrCreateState(sessionKey); + } + + getOrCreateState(sessionKey: string): RunState { + return this.storageService.getOrCreateState(sessionKey); + } + + removeToolCallMapping(toolCallId: string): void { + this.mappingService.removeToolCallMapping(toolCallId); + } + + getStateByToolCallId(toolCallId: string): RunState | undefined { + const sessionKey = this.mappingService.getSessionKey(toolCallId); + if (!sessionKey) return undefined; + return this.storageService.getState(sessionKey); + } + + purgeSessionState(sessionKey: string): void { + // Clean up mappings that point to this session + for (const [toolCallId, sk] of this.mappingService['sessionKeyByToolCallId']) { + if (sk === sessionKey) { + this.mappingService.removeToolCallMapping(toolCallId); + } + } + + this.storageService.purgeSessionState(sessionKey); + } + + hasState(sessionKey: string): boolean { + return this.storageService.hasState(sessionKey); + } + + get stateCount(): number { + return this.storageService.stateCount; + } + + getCleanupCallback(): (action: 'disable' | 'reset' | 'delete' | 'restart') => void { + return this.lifecycleManager.getCleanupCallback(); + } + + reset(): void { + this.mappingService.clear(); + this.storageService.clear(); + } + + addThought(sessionKey: string, thought: ThoughtData): void { + const state = this.getOrCreateState(sessionKey); + + // Use local copy to avoid mutating input + const adjustedInput = { ...thought }; + if (adjustedInput.thoughtNumber > adjustedInput.totalThoughts) { + adjustedInput.totalThoughts = adjustedInput.thoughtNumber; + } + + state.thoughtHistory.push(adjustedInput); + + if (adjustedInput.branchFromThought && adjustedInput.branchId) { + if (!state.branches[adjustedInput.branchId]) { + state.branches[adjustedInput.branchId] = []; + } + state.branches[adjustedInput.branchId].push(adjustedInput); + } + } + + getThoughtHistory(sessionKey: string): ThoughtData[] { + const state = this.getOrCreateState(sessionKey); + return [...state.thoughtHistory]; // Return a copy to prevent external mutations + } + + addBranch(sessionKey: string, branchId: string, thoughts: ThoughtData[]): void { + const state = this.getOrCreateState(sessionKey); + state.branches[branchId] = [...thoughts]; + } + + getBranchIds(sessionKey: string): string[] { + const state = this.getOrCreateState(sessionKey); + return Object.keys(state.branches); + } +} \ No newline at end of file diff --git a/src/hooks.ts b/src/hooks.ts index e2fca42..77c7b55 100644 --- a/src/hooks.ts +++ b/src/hooks.ts @@ -11,7 +11,7 @@ import type { } from "openclaw/plugin-sdk/types"; import { logger } from "../api.js"; import { resolveConfig, type SequentialThinkingConfig } from "./config.js"; -import { SessionStateManager } from "./state.js"; +import { StateManagement } from "./state.js"; import { PREFER_SEQUENTIAL_THINKING_CONTEXT } from "./tool-metadata.js"; type HookConfigContext = { @@ -19,7 +19,7 @@ type HookConfigContext = { }; type HookHandlerDeps = { - manager: SessionStateManager; + manager: StateManagement; registrationConfig: SequentialThinkingConfig; toolName: string; }; @@ -121,7 +121,7 @@ function resolveHookConfig( } function purgeSessionState( - manager: SessionStateManager, + manager: StateManagement, sessionKey: string | undefined, hookName: string, toolName: string, @@ -136,7 +136,7 @@ function purgeSessionState( function createSessionPurgeHandler< TEvent, TContext extends { sessionKey?: string }, ->(manager: SessionStateManager, toolName: string, hookName: string) { +>(manager: StateManagement, toolName: string, hookName: string) { return async (_event: TEvent, ctx: TContext): Promise => { purgeSessionState(manager, ctx.sessionKey, hookName, toolName); }; diff --git a/src/plugin.ts b/src/plugin.ts index b16d91c..1beb819 100644 --- a/src/plugin.ts +++ b/src/plugin.ts @@ -2,7 +2,7 @@ import { type OpenClawPluginApi } from "../api.js"; import { resolveConfig } from "./config.js"; import { createHookHandlers } from "./hooks.js"; import { TOOL_PARAMETER_SCHEMA } from "./schema.js"; -import { SessionStateManager } from "./state.js"; +import { StateManagement, SessionStateManager } from "./state.js"; import { TOOL_DESCRIPTION } from "./tool-metadata.js"; import { SequentialThinkingTool, @@ -81,7 +81,7 @@ export function registerSequentialThinkingPlugin(api: OpenClawPluginApi): void { } function createSessionExtension( - manager: SessionStateManager, + manager: StateManagement, ): SessionExtensionRegistration { return { namespace: TOOL_NAME, diff --git a/src/state-interface.ts b/src/state-interface.ts new file mode 100644 index 0000000..5689357 --- /dev/null +++ b/src/state-interface.ts @@ -0,0 +1,76 @@ +import { RunState, ThoughtData } from './tool.js'; + +/** + * Interface for managing session state in the sequential thinking plugin + */ +export interface StateManagement { + /** + * Register a tool call with a session key + */ + registerToolCall(sessionKey: string, toolCallId: string): void; + + /** + * Get or create state for a session + */ + getOrCreateState(sessionKey: string): RunState; + + /** + * Remove tool call mapping + */ + removeToolCallMapping(toolCallId: string): void; + + /** + * Get state by tool call ID + */ + getStateByToolCallId(toolCallId: string): RunState | undefined; + + /** + * Purge all state for a session + */ + purgeSessionState(sessionKey: string): void; + + /** + * Check if state exists for a session + */ + hasState(sessionKey: string): boolean; + + /** + * Get the number of active states + */ + readonly stateCount: number; + + /** + * Get cleanup callback for session extension + */ + getCleanupCallback(): (action: 'disable' | 'reset' | 'delete' | 'restart') => void; + + /** + * Reset all state + */ + reset(): void; +} + +/** + * Additional interface for state operations + */ +export interface StateOperations { + /** + * Update state with a new thought + */ + addThought(sessionKey: string, thought: ThoughtData): void; + + /** + * Get current thought history + */ + getThoughtHistory(sessionKey: string): ThoughtData[]; + + /** + * Add a branch to the state + */ + addBranch(sessionKey: string, branchId: string, thoughts: ThoughtData[]): void; + + /** + * Get all branch IDs for a session + */ + getBranchIds(sessionKey: string): string[]; +} \ No newline at end of file diff --git a/src/state.ts b/src/state.ts index 2dfa3b5..71b6285 100644 --- a/src/state.ts +++ b/src/state.ts @@ -1,68 +1,13 @@ import { RunState } from "./tool.js"; +import { StateManagement } from "./state-interface.js"; +import { CompositeStateManagement } from "./composite-state-management.js"; -export class SessionStateManager { - private sessionKeyByToolCallId: Map; - private stateBySessionKey: Map; +// Export the interface for use in other modules +export type { StateManagement }; - constructor() { - this.sessionKeyByToolCallId = new Map(); - this.stateBySessionKey = new Map(); - } +// Export the concrete implementation for backward compatibility +export { CompositeStateManagement }; - registerToolCall(sessionKey: string, toolCallId: string): void { - this.sessionKeyByToolCallId.set(toolCallId, sessionKey); - this.getOrCreateState(sessionKey); - } - - getOrCreateState(sessionKey: string): RunState { - if (!this.stateBySessionKey.has(sessionKey)) { - this.stateBySessionKey.set(sessionKey, { - thoughtHistory: [], - branches: {}, - }); - } - return this.stateBySessionKey.get(sessionKey)!; - } - - removeToolCallMapping(toolCallId: string): void { - this.sessionKeyByToolCallId.delete(toolCallId); - } - - getStateByToolCallId(toolCallId: string): RunState | undefined { - const sessionKey = this.sessionKeyByToolCallId.get(toolCallId); - if (!sessionKey) return undefined; - return this.stateBySessionKey.get(sessionKey); - } - - purgeSessionState(sessionKey: string): void { - this.stateBySessionKey.delete(sessionKey); - for (const [toolCallId, sk] of this.sessionKeyByToolCallId) { - if (sk === sessionKey) { - this.sessionKeyByToolCallId.delete(toolCallId); - } - } - } - - hasState(sessionKey: string): boolean { - return this.stateBySessionKey.has(sessionKey); - } - - get stateCount(): number { - return this.stateBySessionKey.size; - } - - getCleanupCallback(): ( - action: "disable" | "reset" | "delete" | "restart", - ) => void { - return () => { - for (const [key] of this.stateBySessionKey) { - this.purgeSessionState(key); - } - }; - } - - reset(): void { - this.sessionKeyByToolCallId.clear(); - this.stateBySessionKey.clear(); - } -} +// For backward compatibility, we alias the new class as the old name +// TODO: Consider renaming all references to use CompositeStateManagement instead +export class SessionStateManager extends CompositeStateManagement {} From 96a8eede8346d6bc14bb5d6295365907c81dce3a Mon Sep 17 00:00:00 2001 From: Ani Date: Sun, 21 Jun 2026 02:13:49 +0000 Subject: [PATCH 2/4] fix: resolve encapsulation violation and memory leak in state management - Add public removeMappingsForSession method to SessionMappingService - Update purgeSessionState to use public API instead of private property access - Fix LifecycleManager to properly clean up both storage and mapping services - Add @deprecated tag to SessionStateManager class - Enhance cleanup callback test to verify tool call mapping removal --- src/composite-state-management.test.ts | 9 ++++---- src/composite-state-management.ts | 29 +++++++++++++++++++------- src/state.ts | 4 +++- 3 files changed, 30 insertions(+), 12 deletions(-) diff --git a/src/composite-state-management.test.ts b/src/composite-state-management.test.ts index 392c2f8..b16aa30 100644 --- a/src/composite-state-management.test.ts +++ b/src/composite-state-management.test.ts @@ -110,10 +110,8 @@ describe("CompositeStateManagement", () => { }); it("should provide cleanup callback", () => { - const cleanupFn = manager.getCleanupCallback(); - expect(typeof cleanupFn).toBe("function"); - - // Add some state + // Add some state and tool call mappings + manager.registerToolCall("session1", "tool1"); const thought: ThoughtData = { thought: "Test thought", thoughtNumber: 1, @@ -122,11 +120,14 @@ describe("CompositeStateManagement", () => { }; manager.addThought("session1", thought); expect(manager.stateCount).toBe(1); + expect(manager.getStateByToolCallId("tool1")).toBeDefined(); // Execute cleanup + const cleanupFn = manager.getCleanupCallback(); cleanupFn("reset"); // State should be cleared expect(manager.stateCount).toBe(0); + expect(manager.getStateByToolCallId("tool1")).toBeUndefined(); }); }); \ No newline at end of file diff --git a/src/composite-state-management.ts b/src/composite-state-management.ts index 011e66b..850e217 100644 --- a/src/composite-state-management.ts +++ b/src/composite-state-management.ts @@ -23,9 +23,25 @@ class SessionMappingService { this.sessionKeyByToolCallId.delete(toolCallId); } + removeMappingsForSession(sessionKey: string): void { + const toRemove: string[] = []; + for (const [toolCallId, sk] of this.sessionKeyByToolCallId) { + if (sk === sessionKey) { + toRemove.push(toolCallId); + } + } + for (const toolCallId of toRemove) { + this.sessionKeyByToolCallId.delete(toolCallId); + } + } + clear(): void { this.sessionKeyByToolCallId.clear(); } + + getAllMappings(): Map { + return new Map(this.sessionKeyByToolCallId); + } } /** @@ -82,15 +98,18 @@ class StateStorageService { */ class LifecycleManager { private storageService: StateStorageService; + private mappingService: SessionMappingService; - constructor(storageService: StateStorageService) { + constructor(storageService: StateStorageService, mappingService: SessionMappingService) { this.storageService = storageService; + this.mappingService = mappingService; } getCleanupCallback(): (action: 'disable' | 'reset' | 'delete' | 'restart') => void { return (action: 'disable' | 'reset' | 'delete' | 'restart') => { for (const key of this.storageService.getAllSessionKeys()) { this.storageService.purgeSessionState(key); + this.mappingService.removeMappingsForSession(key); } }; } @@ -107,7 +126,7 @@ export class CompositeStateManagement implements StateManagement, StateOperation constructor() { this.mappingService = new SessionMappingService(); this.storageService = new StateStorageService(); - this.lifecycleManager = new LifecycleManager(this.storageService); + this.lifecycleManager = new LifecycleManager(this.storageService, this.mappingService); } registerToolCall(sessionKey: string, toolCallId: string): void { @@ -132,11 +151,7 @@ export class CompositeStateManagement implements StateManagement, StateOperation purgeSessionState(sessionKey: string): void { // Clean up mappings that point to this session - for (const [toolCallId, sk] of this.mappingService['sessionKeyByToolCallId']) { - if (sk === sessionKey) { - this.mappingService.removeToolCallMapping(toolCallId); - } - } + this.mappingService.removeMappingsForSession(sessionKey); this.storageService.purgeSessionState(sessionKey); } diff --git a/src/state.ts b/src/state.ts index 71b6285..1977ae4 100644 --- a/src/state.ts +++ b/src/state.ts @@ -9,5 +9,7 @@ export type { StateManagement }; export { CompositeStateManagement }; // For backward compatibility, we alias the new class as the old name -// TODO: Consider renaming all references to use CompositeStateManagement instead +/** + * @deprecated Use CompositeStateManagement or StateManagement interface instead + */ export class SessionStateManager extends CompositeStateManagement {} From 801f64dee1f914ae7e5355d2a4bfb7cc6cef078f Mon Sep 17 00:00:00 2001 From: Ani Date: Sun, 21 Jun 2026 02:23:39 +0000 Subject: [PATCH 3/4] refactor: optimize state management implementation - Simplify removeMappingsForSession to delete directly while iterating - Optimize LifecycleManager cleanup callback to use clear() methods - Clarify getThoughtHistory uses shallow copy in comment - Add tests for reset() method and getThoughtHistory behavior --- src/composite-state-management.test.ts | 60 ++++++++++++++++++++++++++ src/composite-state-management.ts | 14 ++---- 2 files changed, 64 insertions(+), 10 deletions(-) diff --git a/src/composite-state-management.test.ts b/src/composite-state-management.test.ts index b16aa30..265ae8c 100644 --- a/src/composite-state-management.test.ts +++ b/src/composite-state-management.test.ts @@ -130,4 +130,64 @@ describe("CompositeStateManagement", () => { expect(manager.stateCount).toBe(0); expect(manager.getStateByToolCallId("tool1")).toBeUndefined(); }); + + it("should reset all state correctly", () => { + // Add some state and tool call mappings + manager.registerToolCall("session1", "tool1"); + manager.registerToolCall("session2", "tool2"); + const thought1: ThoughtData = { + thought: "Test thought 1", + thoughtNumber: 1, + totalThoughts: 1, + nextThoughtNeeded: false, + }; + const thought2: ThoughtData = { + thought: "Test thought 2", + thoughtNumber: 1, + totalThoughts: 1, + nextThoughtNeeded: false, + }; + manager.addThought("session1", thought1); + manager.addThought("session2", thought2); + + // Verify state exists + expect(manager.stateCount).toBe(2); + expect(manager.getStateByToolCallId("tool1")).toBeDefined(); + expect(manager.getStateByToolCallId("tool2")).toBeDefined(); + + // Execute reset + manager.reset(); + + // All state should be cleared + expect(manager.stateCount).toBe(0); + expect(manager.getStateByToolCallId("tool1")).toBeUndefined(); + expect(manager.getStateByToolCallId("tool2")).toBeUndefined(); + }); + + it("should return a shallow copy of thought history that prevents array mutations", () => { + const thought: ThoughtData = { + thought: "Test thought", + thoughtNumber: 1, + totalThoughts: 1, + nextThoughtNeeded: false, + }; + manager.addThought("session1", thought); + + const history = manager.getThoughtHistory("session1"); + expect(history).toHaveLength(1); + expect(history[0].thought).toBe("Test thought"); + + // Modifying the returned array shouldn't affect internal state + history.push({ + thought: "Added externally", + thoughtNumber: 2, + totalThoughts: 2, + nextThoughtNeeded: false, + }); + + // The internal state should remain unchanged + const internalHistory = manager.getThoughtHistory("session1"); + expect(internalHistory).toHaveLength(1); + expect(internalHistory[0].thought).toBe("Test thought"); + }); }); \ No newline at end of file diff --git a/src/composite-state-management.ts b/src/composite-state-management.ts index 850e217..58abdf5 100644 --- a/src/composite-state-management.ts +++ b/src/composite-state-management.ts @@ -24,15 +24,11 @@ class SessionMappingService { } removeMappingsForSession(sessionKey: string): void { - const toRemove: string[] = []; for (const [toolCallId, sk] of this.sessionKeyByToolCallId) { if (sk === sessionKey) { - toRemove.push(toolCallId); + this.sessionKeyByToolCallId.delete(toolCallId); } } - for (const toolCallId of toRemove) { - this.sessionKeyByToolCallId.delete(toolCallId); - } } clear(): void { @@ -107,10 +103,8 @@ class LifecycleManager { getCleanupCallback(): (action: 'disable' | 'reset' | 'delete' | 'restart') => void { return (action: 'disable' | 'reset' | 'delete' | 'restart') => { - for (const key of this.storageService.getAllSessionKeys()) { - this.storageService.purgeSessionState(key); - this.mappingService.removeMappingsForSession(key); - } + this.storageService.clear(); + this.mappingService.clear(); }; } } @@ -194,7 +188,7 @@ export class CompositeStateManagement implements StateManagement, StateOperation getThoughtHistory(sessionKey: string): ThoughtData[] { const state = this.getOrCreateState(sessionKey); - return [...state.thoughtHistory]; // Return a copy to prevent external mutations + return [...state.thoughtHistory]; // Return a shallow copy to prevent external array mutations (but not mutations of individual ThoughtData properties) } addBranch(sessionKey: string, branchId: string, thoughts: ThoughtData[]): void { From 007072203dc80b5f91a88865c6f672774453277b Mon Sep 17 00:00:00 2001 From: Ani Date: Sun, 21 Jun 2026 03:14:24 +0000 Subject: [PATCH 4/4] refactor: apply final optimization suggestions - Remove unused RunState import from state.ts - Use import type for StateManagement in hooks.ts - Re-export StateOperations interface - Update documentation to include new files --- AGENTS.md | 2 ++ README.md | 2 ++ src/hooks.ts | 2 +- src/state.ts | 3 ++- 4 files changed, 7 insertions(+), 2 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 91bf277..bf526af 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -19,6 +19,8 @@ src/ ├── tool-metadata.ts # Tool description and prompt content ├── schema.ts # Tool parameter validation schema ├── state.ts # State management interface and composite implementation +├── state-interface.ts # StateManagement and StateOperations interfaces +├── composite-state-management.ts # CompositeStateManagement implementation ├── config.ts # Plugin configuration processing ├── plugin.test.ts # Plugin integration tests ├── state.test.ts # State management unit tests diff --git a/README.md b/README.md index 0453600..3c6f1e7 100644 --- a/README.md +++ b/README.md @@ -27,6 +27,8 @@ index.ts ├─ tool.ts → SequentialThinkingTool class ├─ schema.ts → TOOL_PARAMETER_SCHEMA ├─ state.ts → StateManagement interface and CompositeStateManagement implementation + ├─ state-interface.ts → StateManagement and StateOperations interfaces + └─ composite-state-management.ts → CompositeStateManagement class implementation ├─ tool-metadata.ts → TOOL_DESCRIPTION & PREFER_SEQUENTIAL_THINKING_CONTEXT └─ hooks.ts → createHookHandlers() ├─ state.ts (for session management) diff --git a/src/hooks.ts b/src/hooks.ts index 77c7b55..38252e4 100644 --- a/src/hooks.ts +++ b/src/hooks.ts @@ -11,7 +11,7 @@ import type { } from "openclaw/plugin-sdk/types"; import { logger } from "../api.js"; import { resolveConfig, type SequentialThinkingConfig } from "./config.js"; -import { StateManagement } from "./state.js"; +import type { StateManagement } from "./state.js"; import { PREFER_SEQUENTIAL_THINKING_CONTEXT } from "./tool-metadata.js"; type HookConfigContext = { diff --git a/src/state.ts b/src/state.ts index 1977ae4..9e62792 100644 --- a/src/state.ts +++ b/src/state.ts @@ -1,9 +1,10 @@ -import { RunState } from "./tool.js"; import { StateManagement } from "./state-interface.js"; +import { StateOperations } from "./state-interface.js"; import { CompositeStateManagement } from "./composite-state-management.js"; // Export the interface for use in other modules export type { StateManagement }; +export type { StateOperations }; // Export the concrete implementation for backward compatibility export { CompositeStateManagement };