diff --git a/pkg/providers/types.go b/pkg/providers/types.go deleted file mode 100644 index 5433a6607..000000000 --- a/pkg/providers/types.go +++ /dev/null @@ -1,30 +0,0 @@ -package providers - -import ( - "context" - - "github.com/ZanzyTHEbar/dragonscale/pkg/messages" -) - -// Type aliases: canonical types now live in pkg/messages. -// These aliases maintain backward compatibility during migration. -type ToolCall = messages.ToolCall -type FunctionCall = messages.FunctionCall -type UsageInfo = messages.UsageInfo -type Message = messages.Message -type ToolDefinition = messages.ToolDefinition -type ToolFunctionDefinition = messages.ToolFunctionDefinition - -// LLMResponse is the response from an LLM provider API call. -type LLMResponse struct { - Content string `json:"content"` - ToolCalls []ToolCall `json:"tool_calls,omitempty"` - FinishReason string `json:"finish_reason"` - Usage *UsageInfo `json:"usage,omitempty"` -} - -// LLMProvider is the interface for LLM provider implementations. -type LLMProvider interface { - Chat(ctx context.Context, messages []Message, tools []ToolDefinition, model string, options map[string]interface{}) (*LLMResponse, error) - GetDefaultModel() string -} diff --git a/pkg/state/state.go b/pkg/state/state.go index 238f1e693..42f09fe0a 100644 --- a/pkg/state/state.go +++ b/pkg/state/state.go @@ -3,9 +3,6 @@ package state import ( "context" "fmt" - "log" - "os" - "path/filepath" "sync" "time" @@ -40,18 +37,16 @@ func WithDelegate(del memory.MemoryDelegate) Option { return func(m *Manager) { m.delegate = del } } -// Manager manages persistent state with atomic saves. -// When a delegate is present, state persists through agent_kv. -// Otherwise, it falls back to file-based atomic JSON writes. +// Manager manages persistent state through the delegate (agent_kv). type Manager struct { workspace string state *State mu sync.RWMutex - stateFile string delegate memory.MemoryDelegate } -// NewManager creates a new state manager for the given workspace. +// NewManager creates a new state manager with the provided delegate. +// The delegate is required for state persistence. func NewManager(workspace string, opts ...Option) *Manager { sm := &Manager{ workspace: workspace, @@ -64,25 +59,6 @@ func NewManager(workspace string, opts ...Option) *Manager { if sm.delegate != nil { sm.loadFromDelegate() - return sm - } - - stateDir := filepath.Join(workspace, "state") - stateFile := filepath.Join(stateDir, "state.json") - oldStateFile := filepath.Join(workspace, "state.json") - - os.MkdirAll(stateDir, 0755) - sm.stateFile = stateFile - - if _, err := os.Stat(stateFile); os.IsNotExist(err) { - if data, err := os.ReadFile(oldStateFile); err == nil { - if err := jsonv2.Unmarshal(data, sm.state); err == nil { - sm.saveAtomic() - log.Printf("[INFO] state: migrated state from %s to %s", oldStateFile, stateFile) - } - } - } else { - sm.load() } return sm @@ -92,24 +68,8 @@ func (sm *Manager) loadFromDelegate() { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() - // Prefer single-key format (state:json) if v, err := sm.delegate.GetKV(ctx, kvAgentID, "state:json"); err == nil && v != "" { - if err := jsonv2.Unmarshal([]byte(v), sm.state); err == nil { - return - } - } - - // Fallback: legacy 3-key format - if v, err := sm.delegate.GetKV(ctx, kvAgentID, "state:last_channel"); err == nil && v != "" { - sm.state.LastChannel = v - } - if v, err := sm.delegate.GetKV(ctx, kvAgentID, "state:last_chat_id"); err == nil && v != "" { - sm.state.LastChatID = v - } - if v, err := sm.delegate.GetKV(ctx, kvAgentID, "state:timestamp"); err == nil && v != "" { - if t, err := time.Parse(time.RFC3339Nano, v); err == nil { - sm.state.Timestamp = t - } + _ = jsonv2.Unmarshal([]byte(v), sm.state) } } @@ -147,13 +107,13 @@ func (sm *Manager) SetChannelAndChatID(ctx context.Context, channel, chatID stri return sm.persist(ctx) } -// persist writes the current state to the delegate (KV) or file. +// persist writes the current state to the delegate (KV) if available. // Must be called with the lock held. func (sm *Manager) persist(ctx context.Context) error { - if sm.delegate != nil { - return sm.persistToDelegate(ctx) + if sm.delegate == nil { + return nil } - return sm.saveAtomic() + return sm.persistToDelegate(ctx) } func (sm *Manager) persistToDelegate(ctx context.Context) error { @@ -187,53 +147,3 @@ func (sm *Manager) GetTimestamp() time.Time { defer sm.mu.RUnlock() return sm.state.Timestamp } - -// saveAtomic performs an atomic save using temp file + rename. -// This ensures that the state file is never corrupted: -// 1. Write to a temp file -// 2. Rename temp file to target (atomic on POSIX systems) -// 3. If rename fails, cleanup the temp file -// -// Must be called with the lock held. -func (sm *Manager) saveAtomic() error { - // Create temp file in the same directory as the target - tempFile := sm.stateFile + ".tmp" - - // Marshal state to JSON - data, err := jsonv2.Marshal(sm.state, jsontext.WithIndent(" ")) - if err != nil { - return fmt.Errorf("failed to marshal state: %w", err) - } - - // Write to temp file - if err := os.WriteFile(tempFile, data, 0644); err != nil { - return fmt.Errorf("failed to write temp file: %w", err) - } - - // Atomic rename from temp to target - if err := os.Rename(tempFile, sm.stateFile); err != nil { - // Cleanup temp file if rename fails - os.Remove(tempFile) - return fmt.Errorf("failed to rename temp file: %w", err) - } - - return nil -} - -// load loads the state from disk. -func (sm *Manager) load() error { - data, err := os.ReadFile(sm.stateFile) - if err != nil { - // File doesn't exist yet, that's OK - if os.IsNotExist(err) { - return nil - } - return fmt.Errorf("failed to read state file: %w", err) - } - - if err := jsonv2.Unmarshal(data, sm.state); err != nil { - return fmt.Errorf("failed to unmarshal state: %w", err) - } - - return nil -} diff --git a/pkg/state/state_test.go b/pkg/state/state_test.go index 0d512ab0a..929cbdc25 100644 --- a/pkg/state/state_test.go +++ b/pkg/state/state_test.go @@ -1,34 +1,26 @@ package state import ( - "fmt" - "os" - "path/filepath" + "context" "testing" - - jsonv2 "github.com/go-json-experiment/json" ) -func TestAtomicSave(t *testing.T) { - t.Parallel( - // Create temp workspace - ) +// Tests for state manager with in-memory state (no delegate). +// Persistence is tested via integration tests with a real delegate. - tmpDir, err := os.MkdirTemp("", "state-test-*") - if err != nil { - t.Fatalf("Failed to create temp dir: %v", err) - } - defer os.RemoveAll(tmpDir) +func TestSetLastChannel(t *testing.T) { + t.Parallel() - sm := NewManager(tmpDir) + sm := NewManager("/tmp/test") // Test SetLastChannel - err = sm.SetLastChannel(t.Context(), "test-channel") + err := sm.SetLastChannel(context.Background(), "test-channel") + // Without a delegate, persist returns nil but state is updated in memory if err != nil { t.Fatalf("SetLastChannel failed: %v", err) } - // Verify the channel was saved + // Verify the channel was saved in memory lastChannel := sm.GetLastChannel() if lastChannel != "test-channel" { t.Errorf("Expected channel 'test-channel', got '%s'", lastChannel) @@ -38,37 +30,20 @@ func TestAtomicSave(t *testing.T) { if sm.GetTimestamp().IsZero() { t.Error("Expected timestamp to be updated") } - - // Verify state file exists - stateFile := filepath.Join(tmpDir, "state", "state.json") - if _, err := os.Stat(stateFile); os.IsNotExist(err) { - t.Error("Expected state file to exist") - } - - // Create a new manager to verify persistence - sm2 := NewManager(tmpDir) - if sm2.GetLastChannel() != "test-channel" { - t.Errorf("Expected persistent channel 'test-channel', got '%s'", sm2.GetLastChannel()) - } } func TestSetLastChatID(t *testing.T) { t.Parallel() - tmpDir, err := os.MkdirTemp("", "state-test-*") - if err != nil { - t.Fatalf("Failed to create temp dir: %v", err) - } - defer os.RemoveAll(tmpDir) - sm := NewManager(tmpDir) + sm := NewManager("/tmp/test") // Test SetLastChatID - err = sm.SetLastChatID(t.Context(), "test-chat-id") + err := sm.SetLastChatID(context.Background(), "test-chat-id") if err != nil { t.Fatalf("SetLastChatID failed: %v", err) } - // Verify the chat ID was saved + // Verify the chat ID was saved in memory lastChatID := sm.GetLastChatID() if lastChatID != "test-chat-id" { t.Errorf("Expected chat ID 'test-chat-id', got '%s'", lastChatID) @@ -78,137 +53,33 @@ func TestSetLastChatID(t *testing.T) { if sm.GetTimestamp().IsZero() { t.Error("Expected timestamp to be updated") } - - // Create a new manager to verify persistence - sm2 := NewManager(tmpDir) - if sm2.GetLastChatID() != "test-chat-id" { - t.Errorf("Expected persistent chat ID 'test-chat-id', got '%s'", sm2.GetLastChatID()) - } } -func TestAtomicity_NoCorruptionOnInterrupt(t *testing.T) { +func TestSetChannelAndChatID(t *testing.T) { t.Parallel() - tmpDir, err := os.MkdirTemp("", "state-test-*") + + sm := NewManager("/tmp/test") + + // Test setting both channel and chat ID atomically + err := sm.SetChannelAndChatID(context.Background(), "channel-1", "chat-1") if err != nil { - t.Fatalf("Failed to create temp dir: %v", err) - } - defer os.RemoveAll(tmpDir) - - sm := NewManager(tmpDir) - - // Write initial state - err = sm.SetLastChannel(t.Context(), "initial-channel") - if err != nil { - t.Fatalf("SetLastChannel failed: %v", err) + t.Fatalf("SetChannelAndChatID failed: %v", err) } - // Simulate a crash scenario by manually creating a corrupted temp file - tempFile := filepath.Join(tmpDir, "state", "state.json.tmp") - err = os.WriteFile(tempFile, []byte("corrupted data"), 0644) - if err != nil { - t.Fatalf("Failed to create temp file: %v", err) + // Verify both were saved in memory + if sm.GetLastChannel() != "channel-1" { + t.Errorf("Expected channel 'channel-1', got '%s'", sm.GetLastChannel()) } - // Verify that the original state is still intact - lastChannel := sm.GetLastChannel() - if lastChannel != "initial-channel" { - t.Errorf("Expected channel 'initial-channel' after corrupted temp file, got '%s'", lastChannel) - } - - // Clean up the temp file manually - os.Remove(tempFile) - - // Now do a proper save - err = sm.SetLastChannel(t.Context(), "new-channel") - if err != nil { - t.Fatalf("SetLastChannel failed: %v", err) - } - - // Verify the new state was saved - if sm.GetLastChannel() != "new-channel" { - t.Errorf("Expected channel 'new-channel', got '%s'", sm.GetLastChannel()) - } -} - -func TestConcurrentAccess(t *testing.T) { - t.Parallel() - tmpDir, err := os.MkdirTemp("", "state-test-*") - if err != nil { - t.Fatalf("Failed to create temp dir: %v", err) - } - defer os.RemoveAll(tmpDir) - - sm := NewManager(tmpDir) - - // Test concurrent writes - done := make(chan bool, 10) - for i := 0; i < 10; i++ { - go func(idx int) { - channel := fmt.Sprintf("channel-%d", idx) - sm.SetLastChannel(t.Context(), channel) - done <- true - }(i) - } - - // Wait for all goroutines to complete - for i := 0; i < 10; i++ { - <-done - } - - // Verify the final state is consistent - lastChannel := sm.GetLastChannel() - if lastChannel == "" { - t.Error("Expected non-empty channel after concurrent writes") - } - - // Verify state file is valid JSON - stateFile := filepath.Join(tmpDir, "state", "state.json") - data, err := os.ReadFile(stateFile) - if err != nil { - t.Fatalf("Failed to read state file: %v", err) - } - - var state State - if err := jsonv2.Unmarshal(data, &state); err != nil { - t.Errorf("State file contains invalid JSON: %v", err) - } -} - -func TestNewManager_ExistingState(t *testing.T) { - t.Parallel() - tmpDir, err := os.MkdirTemp("", "state-test-*") - if err != nil { - t.Fatalf("Failed to create temp dir: %v", err) - } - defer os.RemoveAll(tmpDir) - - // Create initial state - sm1 := NewManager(tmpDir) - sm1.SetLastChannel(t.Context(), "existing-channel") - sm1.SetLastChatID(t.Context(), "existing-chat-id") - - // Create new manager with same workspace - sm2 := NewManager(tmpDir) - - // Verify state was loaded - if sm2.GetLastChannel() != "existing-channel" { - t.Errorf("Expected channel 'existing-channel', got '%s'", sm2.GetLastChannel()) - } - - if sm2.GetLastChatID() != "existing-chat-id" { - t.Errorf("Expected chat ID 'existing-chat-id', got '%s'", sm2.GetLastChatID()) + if sm.GetLastChatID() != "chat-1" { + t.Errorf("Expected chat ID 'chat-1', got '%s'", sm.GetLastChatID()) } } func TestNewManager_EmptyWorkspace(t *testing.T) { t.Parallel() - tmpDir, err := os.MkdirTemp("", "state-test-*") - if err != nil { - t.Fatalf("Failed to create temp dir: %v", err) - } - defer os.RemoveAll(tmpDir) - sm := NewManager(tmpDir) + sm := NewManager("/tmp/test") // Verify default state if sm.GetLastChannel() != "" { @@ -223,3 +94,21 @@ func TestNewManager_EmptyWorkspace(t *testing.T) { t.Error("Expected zero timestamp for new state") } } + +func TestStateStruct(t *testing.T) { + t.Parallel() + + // Test that State struct fields work correctly + state := &State{ + LastChannel: "test-channel", + LastChatID: "test-chat-id", + } + + if state.LastChannel != "test-channel" { + t.Errorf("Expected LastChannel 'test-channel', got '%s'", state.LastChannel) + } + + if state.LastChatID != "test-chat-id" { + t.Errorf("Expected LastChatID 'test-chat-id', got '%s'", state.LastChatID) + } +} diff --git a/pkg/tools/shell.go b/pkg/tools/shell.go index ef5ff16b1..23ec923a9 100644 --- a/pkg/tools/shell.go +++ b/pkg/tools/shell.go @@ -41,26 +41,26 @@ const ( ) type ExecTool struct { - workingDir string - timeout time.Duration - denyPatterns []*regexp.Regexp - allowPatterns []*regexp.Regexp - restrictToWorkspace bool - mode ShellMode - workspace string // root workspace path for audit logging + workingDir string + timeout time.Duration + denyPatterns []*regexp.Regexp + allowPatterns []*regexp.Regexp + restrictToSandbox bool + mode ShellMode + workspace string // root workspace path for audit logging } func NewExecTool(workingDir string, restrict bool) *ExecTool { denyPatterns := buildDenyPatterns() return &ExecTool{ - workingDir: workingDir, - timeout: 60 * time.Second, - denyPatterns: denyPatterns, - allowPatterns: nil, - restrictToWorkspace: restrict, - mode: ShellModeDenyList, - workspace: workingDir, + workingDir: workingDir, + timeout: 60 * time.Second, + denyPatterns: denyPatterns, + allowPatterns: nil, + restrictToSandbox: restrict, + mode: ShellModeDenyList, + workspace: workingDir, } } @@ -196,7 +196,7 @@ func (t *ExecTool) Execute(ctx context.Context, args map[string]interface{}) *To cwd := t.workingDir if wd, ok := args["working_dir"].(string); ok && wd != "" { - if t.restrictToWorkspace && t.workspace != "" { + if t.restrictToSandbox && t.workspace != "" { resolved, err := validatePath(wd, t.workspace, true) if err != nil { return ErrorResult(fmt.Sprintf("working_dir blocked: %v", err)) @@ -385,7 +385,7 @@ func (t *ExecTool) guardCommand(command, cwd string) string { } } - if t.restrictToWorkspace { + if t.restrictToSandbox { if strings.Contains(cmd, "..\\") || strings.Contains(cmd, "../") { return "Command blocked by safety guard (path traversal detected)" } @@ -443,13 +443,8 @@ func (t *ExecTool) SetTimeout(timeout time.Duration) { t.timeout = timeout } -// SetRestrictToWorkspace is the legacy name. Use SetRestrictToSandbox for new code. -func (t *ExecTool) SetRestrictToWorkspace(restrict bool) { - t.restrictToWorkspace = restrict -} - func (t *ExecTool) SetRestrictToSandbox(restrict bool) { - t.restrictToWorkspace = restrict + t.restrictToSandbox = restrict } func (t *ExecTool) SetMode(mode ShellMode) { diff --git a/pkg/tools/shell_test.go b/pkg/tools/shell_test.go index 3c1a6b478..4ef83d618 100644 --- a/pkg/tools/shell_test.go +++ b/pkg/tools/shell_test.go @@ -195,12 +195,12 @@ func TestShellTool_OutputTruncation(t *testing.T) { } } -// TestShellTool_RestrictToWorkspace verifies workspace restriction -func TestShellTool_RestrictToWorkspace(t *testing.T) { +// TestShellTool_RestrictToSandbox verifies sandbox restriction +func TestShellTool_RestrictToSandbox(t *testing.T) { t.Parallel() tmpDir := t.TempDir() tool := NewExecTool(tmpDir, false) - tool.SetRestrictToWorkspace(true) + tool.SetRestrictToSandbox(true) ctx := t.Context() args := map[string]interface{}{ @@ -211,7 +211,7 @@ func TestShellTool_RestrictToWorkspace(t *testing.T) { // Path traversal should be blocked if !result.IsError { - t.Errorf("Expected path traversal to be blocked with restrictToWorkspace=true") + t.Errorf("Expected path traversal to be blocked with restrictToSandbox=true") } if !strings.Contains(result.ForLLM, "blocked") && !strings.Contains(result.ForUser, "blocked") { diff --git a/pkg/tools/types.go b/pkg/tools/types.go deleted file mode 100644 index f8205b8bd..000000000 --- a/pkg/tools/types.go +++ /dev/null @@ -1,52 +0,0 @@ -package tools - -import "context" - -type Message struct { - Role string `json:"role"` - Content string `json:"content"` - ToolCalls []ToolCall `json:"tool_calls,omitempty"` - ToolCallID string `json:"tool_call_id,omitempty"` -} - -type ToolCall struct { - ID string `json:"id"` - Type string `json:"type"` - Function *FunctionCall `json:"function,omitempty"` - Name string `json:"name,omitempty"` - Arguments map[string]interface{} `json:"arguments,omitempty"` -} - -type FunctionCall struct { - Name string `json:"name"` - Arguments string `json:"arguments"` -} - -type LLMResponse struct { - Content string `json:"content"` - ToolCalls []ToolCall `json:"tool_calls,omitempty"` - FinishReason string `json:"finish_reason"` - Usage *UsageInfo `json:"usage,omitempty"` -} - -type UsageInfo struct { - PromptTokens int `json:"prompt_tokens"` - CompletionTokens int `json:"completion_tokens"` - TotalTokens int `json:"total_tokens"` -} - -type LLMProvider interface { - Chat(ctx context.Context, messages []Message, tools []ToolDefinition, model string, options map[string]interface{}) (*LLMResponse, error) - GetDefaultModel() string -} - -type ToolDefinition struct { - Type string `json:"type"` - Function ToolFunctionDefinition `json:"function"` -} - -type ToolFunctionDefinition struct { - Name string `json:"name"` - Description string `json:"description"` - Parameters map[string]interface{} `json:"parameters"` -}