opencode/picoclaw
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

refactor(providers): remove types.go

Browse source 
test; remove types.go
This commit is contained in:
ZanzyTHEbar 2026-03-05 14:26:24 +00:00
parent 7f64dd1cc5
commit a600ae7968
6 changed files with 72 additions and 360 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

30
pkg/providers/types.go
View file

@ -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
}

106
pkg/state/state.go
View file

@ -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
}

197
pkg/state/state_test.go
View file

@ -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)
}
}

15
pkg/tools/shell.go
View file

@ -45,7 +45,7 @@ type ExecTool struct {
timeout time.Duration
denyPatterns []*regexp.Regexp
allowPatterns []*regexp.Regexp
restrictToWorkspace bool
restrictToSandbox bool
mode ShellMode
workspace string // root workspace path for audit logging
}
@ -58,7 +58,7 @@ func NewExecTool(workingDir string, restrict bool) *ExecTool {
timeout: 60 * time.Second,
denyPatterns: denyPatterns,
allowPatterns: nil,
restrictToWorkspace: restrict,
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) {

8
pkg/tools/shell_test.go
View file

@ -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") {

52
pkg/tools/types.go
View file

@ -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"`
}