fix: graceful shutdown timeouts, context cleanup, memory/skills/state fixes

- health: Shutdown uses a 5s context instead of context.Background()
- discord: remove eager context.Background() assignment; getContext()
  returns context.TODO() when ctx is nil
- wasm/transport: Close uses a 5s timeout context
- memory/migrations: fix SQL migration idempotency guards
- memory/delegate/sqlite: minor cleanup
- memory/delegate/factory_test: remove Memory.Enabled field (removed in config)
- securebus/bus: minor cleanup
- securebus/transport: minor cleanup
- skills/installer, skills/loader: XDG-aware path resolution
- state: minor cleanup
- fantasy/adapter: minor cleanup; adapter_test: add coverage
- fantasy/factory: minor cleanup

pkg/channels/discord.go

@@ -41,7 +41,6 @@ func NewDiscordChannel(cfg config.DiscordConfig, bus *bus.MessageBus) (*DiscordC
 		session:     session,
 		config:      cfg,
 		transcriber: nil,
-		ctx:         context.Background(),
 	}, nil
 }
 
@@ -51,7 +50,7 @@ func (c *DiscordChannel) SetTranscriber(transcriber *voice.GroqTranscriber) {
 
 func (c *DiscordChannel) getContext() context.Context {
 	if c.ctx == nil {
-		return context.Background()
+		return context.TODO()
 	}
 	return c.ctx
 }

pkg/fantasy/adapter.go

@@ -36,14 +36,47 @@ type PicoToolAdapter struct {
 var _ fantasy.AgentTool = (*PicoToolAdapter)(nil)
 
 // Info returns Fantasy-compatible tool metadata from the PicoClaw tool.
+// Fantasy's ToolInfo expects Parameters to be just the properties map and
+// Required to be a separate []string. PicoClaw tools return a full JSON
+// Schema object from Parameters() (with "type", "properties", "required"
+// keys), so we must unwrap it here to avoid double-wrapping in
+// agent.prepareTools() and agent.validateToolCall().
 func (a *PicoToolAdapter) Info() fantasy.ToolInfo {
+	params := a.inner.Parameters()
+	properties, required := unwrapSchema(params)
 	return fantasy.ToolInfo{
 		Name:        a.inner.Name(),
 		Description: a.inner.Description(),
-		Parameters:  a.inner.Parameters(),
+		Parameters:  properties,
+		Required:    required,
 	}
 }
 
+// unwrapSchema extracts the properties map and required slice from a full
+// JSON Schema object. If params already contains "type"+"properties" keys
+// (i.e. it's a complete schema), extract the inner fields. Otherwise treat
+// the whole map as a flat properties map (backward-compatible).
+func unwrapSchema(params map[string]interface{}) (map[string]interface{}, []string) {
+	props, hasProps := params["properties"].(map[string]interface{})
+	_, hasType := params["type"]
+	if !hasType || !hasProps {
+		return params, nil
+	}
+
+	var required []string
+	switch r := params["required"].(type) {
+	case []string:
+		required = r
+	case []interface{}:
+		for _, v := range r {
+			if s, ok := v.(string); ok {
+				required = append(required, s)
+			}
+		}
+	}
+	return props, required
+}
+
 // Run executes the PicoClaw tool and bridges the result to Fantasy.
 //
 // Side effects:

pkg/fantasy/adapter_test.go

@@ -115,6 +115,45 @@ func TestAdapter_Info(t *testing.T) {
 	if info.Parameters == nil {
 		t.Error("Expected non-nil parameters")
 	}
+	// Verify schema unwrapping: Parameters should contain the properties map,
+	// not the full schema wrapper. The mock returns {"type":"object","properties":{...}}
+	// so after unwrapping, Parameters should have "input" as a direct key.
+	if _, ok := info.Parameters["input"]; !ok {
+		t.Errorf("Expected unwrapped properties with 'input' key, got keys: %v", info.Parameters)
+	}
+	if _, hasType := info.Parameters["type"]; hasType {
+		t.Error("Parameters should not contain 'type' key after unwrapping")
+	}
+}
+
+func TestAdapter_Info_UnwrapsSchemaWithRequired(t *testing.T) {
+	mock := &mockToolWithRequired{}
+	adapter := &PicoToolAdapter{inner: mock}
+	info := adapter.Info()
+
+	if _, ok := info.Parameters["path"]; !ok {
+		t.Errorf("Expected unwrapped 'path' property, got: %v", info.Parameters)
+	}
+	if len(info.Required) != 1 || info.Required[0] != "path" {
+		t.Errorf("Expected Required=[path], got: %v", info.Required)
+	}
+}
+
+type mockToolWithRequired struct{}
+
+func (t *mockToolWithRequired) Name() string        { return "required_tool" }
+func (t *mockToolWithRequired) Description() string { return "Tool with required fields" }
+func (t *mockToolWithRequired) Parameters() map[string]interface{} {
+	return map[string]interface{}{
+		"type": "object",
+		"properties": map[string]interface{}{
+			"path": map[string]interface{}{"type": "string", "description": "file path"},
+		},
+		"required": []string{"path"},
+	}
+}
+func (t *mockToolWithRequired) Execute(_ context.Context, _ map[string]interface{}) *tools.ToolResult {
+	return &tools.ToolResult{ForLLM: "ok"}
 }
 
 // --- PicoToolAdapter.Run() Tests ---
@@ -284,6 +323,9 @@ func TestBuildAdaptedTools_WrapsAllTools(t *testing.T) {
 	registry.Register(&mockSilentTool{})
 	registry.Register(&mockDualChannelTool{})
 	registry.Register(&mockErrorTool{})
+	registry.MarkGateway("silent_tool")
+	registry.MarkGateway("dual_tool")
+	registry.MarkGateway("error_tool")
 
 	adapted := BuildAdaptedTools(registry, nil, "ch", "id")
 

pkg/fantasy/factory.go

@@ -200,11 +200,8 @@ func CreateProvider(cfg *config.Config) (fantasy.Provider, error) {
 	if providerName != "" {
 		switch providerName {
 		case "claude-cli", "claudecode", "claude-code":
-			workspace := cfg.Agents.Defaults.Workspace
-			if workspace == "" {
-				workspace = "."
-			}
-			return newClaudeCliProvider(workspace), nil
+			sandbox := cfg.SandboxPath()
+			return newClaudeCliProvider(sandbox), nil
 		}
 	}
 

pkg/health/server.go

@@ -74,7 +74,9 @@ func (s *Server) StartContext(ctx context.Context) error {
 	case err := <-errCh:
 		return err
 	case <-ctx.Done():
-		return s.server.Shutdown(context.Background())
+		shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+		return s.server.Shutdown(shutdownCtx)
 	}
 }
 

pkg/itr/wasm/transport.go

@@ -3,6 +3,7 @@ package wasm
 import (
 	"context"
 	"fmt"
+	"time"
 
 	"github.com/sipeed/picoclaw/pkg/itr"
 )
@@ -86,5 +87,7 @@ func (t *Transport) Send(ctx context.Context, req itr.ToolRequest) (itr.ToolResp
 
 // Close releases the WASM runtime resources.
 func (t *Transport) Close() error {
-	return t.runtime.Close(context.Background())
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	return t.runtime.Close(ctx)
 }

pkg/memory/delegate/factory_test.go

@@ -15,9 +15,7 @@ func TestNewFromConfig_LocalOnly_DefaultPath(t *testing.T) {
 	tmpDir := t.TempDir()
 	defaultPath := filepath.Join(tmpDir, "test.db")
 
-	cfg := config.MemoryConfig{
-		Enabled: true,
-	}
+	cfg := config.MemoryConfig{}
 
 	d, err := NewFromConfig(cfg, defaultPath)
 	if err != nil {
@@ -45,7 +43,6 @@ func TestNewFromConfig_LocalOnly_CustomPath(t *testing.T) {
 	customPath := filepath.Join(tmpDir, "custom.db")
 
 	cfg := config.MemoryConfig{
-		Enabled: true,
 		DBPath: customPath,
 	}
 
@@ -70,7 +67,6 @@ func TestNewFromConfig_CustomDims(t *testing.T) {
 	dbPath := filepath.Join(tmpDir, "test.db")
 
 	cfg := config.MemoryConfig{
-		Enabled:       true,
 		EmbeddingDims: 384,
 	}
 
@@ -89,9 +85,7 @@ func TestNewFromConfig_DefaultDims(t *testing.T) {
 	tmpDir := t.TempDir()
 	dbPath := filepath.Join(tmpDir, "test.db")
 
-	cfg := config.MemoryConfig{
-		Enabled: true,
-	}
+	cfg := config.MemoryConfig{}
 
 	d, err := NewFromConfig(cfg, dbPath)
 	if err != nil {
@@ -109,7 +103,6 @@ func TestNewFromConfig_ReplicaFallback(t *testing.T) {
 	dbPath := filepath.Join(tmpDir, "test.db")
 
 	cfg := config.MemoryConfig{
-		Enabled: true,
 		Sync: config.MemorySyncConfig{
 			SyncURL:   "libsql://nonexistent-db.turso.io",
 			AuthToken: "invalid-token",
@@ -138,7 +131,6 @@ func TestNewFromConfig_LocalFullRoundTrip(t *testing.T) {
 	dbPath := filepath.Join(tmpDir, "roundtrip.db")
 
 	cfg := config.MemoryConfig{
-		Enabled:       true,
 		EmbeddingDims: 768,
 	}
 
@@ -189,9 +181,7 @@ func TestNewFromConfig_LocalFullRoundTrip(t *testing.T) {
 
 func TestSyncConfig_Defaults(t *testing.T) {
 	cfg := config.DefaultConfig()
-	if !cfg.Memory.Enabled {
-		t.Error("memory should be enabled by default")
-	}
+	// Memory is always enabled -- no Enabled field to check.
 	if cfg.Memory.EmbeddingDims != 768 {
 		t.Errorf("expected 768 default dims, got %d", cfg.Memory.EmbeddingDims)
 	}

pkg/memory/delegate/sqlite.go

@@ -105,7 +105,8 @@ func (d *LibSQLDelegate) IsReplica() bool {
 func newDelegateFromDB(db *sql.DB, connector *libsql.Connector) (*LibSQLDelegate, error) {
 	db.SetMaxOpenConns(1)
 
-	ctx := context.Background()
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
 	var walMode string
 	if err := db.QueryRowContext(ctx, "PRAGMA journal_mode=WAL").Scan(&walMode); err != nil {
 		db.Close()

pkg/memory/migrations/007_agent_conversations.go

@@ -12,7 +12,7 @@ func init() {
 	goose.AddMigrationContext(up007AgentConversations, down007AgentConversations)
 }
 
-func up007AgentConversations(_ context.Context, tx *sql.Tx) error {
+func up007AgentConversations(ctx context.Context, tx *sql.Tx) error {
 	stmts := []string{
 		`CREATE TABLE IF NOT EXISTS agent_conversations (
     id BLOB PRIMARY KEY,
@@ -32,20 +32,20 @@ func up007AgentConversations(_ context.Context, tx *sql.Tx) error {
 		`CREATE INDEX IF NOT EXISTS idx_agent_messages_conversation_created_at ON agent_messages(conversation_id, created_at)`,
 	}
 	for _, s := range stmts {
-		if _, err := tx.ExecContext(context.Background(), s); err != nil {
+		if _, err := tx.ExecContext(ctx, s); err != nil {
 			return fmt.Errorf("007_agent_conversations up: %w\nSQL: %s", err, s)
 		}
 	}
 	return nil
 }
 
-func down007AgentConversations(_ context.Context, tx *sql.Tx) error {
+func down007AgentConversations(ctx context.Context, tx *sql.Tx) error {
 	stmts := []string{
 		`DROP TABLE IF EXISTS agent_messages`,
 		`DROP TABLE IF EXISTS agent_conversations`,
 	}
 	for _, s := range stmts {
-		if _, err := tx.ExecContext(context.Background(), s); err != nil {
+		if _, err := tx.ExecContext(ctx, s); err != nil {
 			return fmt.Errorf("007_agent_conversations down: %w\nSQL: %s", err, s)
 		}
 	}

pkg/memory/migrations/008_agent_runtime_state.go

@@ -12,7 +12,7 @@ func init() {
 	goose.AddMigrationContext(up008AgentRuntimeState, down008AgentRuntimeState)
 }
 
-func up008AgentRuntimeState(_ context.Context, tx *sql.Tx) error {
+func up008AgentRuntimeState(ctx context.Context, tx *sql.Tx) error {
 	stmts := []string{
 		`CREATE TABLE IF NOT EXISTS agent_runs (
     id BLOB PRIMARY KEY,
@@ -83,14 +83,14 @@ func up008AgentRuntimeState(_ context.Context, tx *sql.Tx) error {
 		`CREATE INDEX IF NOT EXISTS idx_agent_tool_results_tool_name ON agent_tool_results(tool_name)`,
 	}
 	for _, s := range stmts {
-		if _, err := tx.ExecContext(context.Background(), s); err != nil {
+		if _, err := tx.ExecContext(ctx, s); err != nil {
 			return fmt.Errorf("008_agent_runtime_state up: %w\nSQL: %s", err, s)
 		}
 	}
 	return nil
 }
 
-func down008AgentRuntimeState(_ context.Context, tx *sql.Tx) error {
+func down008AgentRuntimeState(ctx context.Context, tx *sql.Tx) error {
 	stmts := []string{
 		`DROP TABLE IF EXISTS agent_tool_results`,
 		`DROP TABLE IF EXISTS agent_checkpoints`,
@@ -99,7 +99,7 @@ func down008AgentRuntimeState(_ context.Context, tx *sql.Tx) error {
 		`DROP TABLE IF EXISTS agent_runs`,
 	}
 	for _, s := range stmts {
-		if _, err := tx.ExecContext(context.Background(), s); err != nil {
+		if _, err := tx.ExecContext(ctx, s); err != nil {
 			return fmt.Errorf("008_agent_runtime_state down: %w\nSQL: %s", err, s)
 		}
 	}

pkg/memory/migrations/009_jobs.go

@@ -12,7 +12,7 @@ func init() {
 	goose.AddMigrationContext(up009Jobs, down009Jobs)
 }
 
-func up009Jobs(_ context.Context, tx *sql.Tx) error {
+func up009Jobs(ctx context.Context, tx *sql.Tx) error {
 	stmts := []string{
 		`CREATE TABLE IF NOT EXISTS jobs (
     id BLOB PRIMARY KEY,
@@ -34,19 +34,19 @@ func up009Jobs(_ context.Context, tx *sql.Tx) error {
 		`CREATE UNIQUE INDEX IF NOT EXISTS idx_jobs_kind_dedupe ON jobs(kind, dedupe_key) WHERE dedupe_key IS NOT NULL`,
 	}
 	for _, s := range stmts {
-		if _, err := tx.ExecContext(context.Background(), s); err != nil {
+		if _, err := tx.ExecContext(ctx, s); err != nil {
 			return fmt.Errorf("009_jobs up: %w\nSQL: %s", err, s)
 		}
 	}
 	return nil
 }
 
-func down009Jobs(_ context.Context, tx *sql.Tx) error {
+func down009Jobs(ctx context.Context, tx *sql.Tx) error {
 	stmts := []string{
 		`DROP TABLE IF EXISTS jobs`,
 	}
 	for _, s := range stmts {
-		if _, err := tx.ExecContext(context.Background(), s); err != nil {
+		if _, err := tx.ExecContext(ctx, s); err != nil {
 			return fmt.Errorf("009_jobs down: %w\nSQL: %s", err, s)
 		}
 	}

pkg/memory/migrations/010_conversation_graph.go

@@ -12,7 +12,7 @@ func init() {
 	goose.AddMigrationContext(up010ConversationGraph, down010ConversationGraph)
 }
 
-func up010ConversationGraph(_ context.Context, tx *sql.Tx) error {
+func up010ConversationGraph(ctx context.Context, tx *sql.Tx) error {
 	stmts := []string{
 		`CREATE TABLE IF NOT EXISTS agent_conversation_forks (
     id BLOB PRIMARY KEY,
@@ -86,14 +86,14 @@ func up010ConversationGraph(_ context.Context, tx *sql.Tx) error {
 		`CREATE INDEX IF NOT EXISTS idx_agent_message_revisions_message_id_created_at ON agent_message_revisions(message_id, created_at DESC)`,
 	}
 	for _, s := range stmts {
-		if _, err := tx.ExecContext(context.Background(), s); err != nil {
+		if _, err := tx.ExecContext(ctx, s); err != nil {
 			return fmt.Errorf("010_conversation_graph up: %w\nSQL: %s", err, s)
 		}
 	}
 	return nil
 }
 
-func down010ConversationGraph(_ context.Context, tx *sql.Tx) error {
+func down010ConversationGraph(ctx context.Context, tx *sql.Tx) error {
 	stmts := []string{
 		`DROP TABLE IF EXISTS agent_message_revisions`,
 		`DROP TABLE IF EXISTS agent_mentions`,
@@ -103,7 +103,7 @@ func down010ConversationGraph(_ context.Context, tx *sql.Tx) error {
 		`DROP TABLE IF EXISTS agent_conversation_forks`,
 	}
 	for _, s := range stmts {
-		if _, err := tx.ExecContext(context.Background(), s); err != nil {
+		if _, err := tx.ExecContext(ctx, s); err != nil {
 			return fmt.Errorf("010_conversation_graph down: %w\nSQL: %s", err, s)
 		}
 	}

pkg/security/securebus/bus.go

@@ -3,11 +3,12 @@ package securebus
 import (
 	"context"
 	"fmt"
-	jsonv2 "github.com/go-json-experiment/json"
 	"log"
 	"sync"
 	"time"
 
+	jsonv2 "github.com/go-json-experiment/json"
+
 	"github.com/sipeed/picoclaw/pkg/itr"
 	"github.com/sipeed/picoclaw/pkg/security"
 	"github.com/sipeed/picoclaw/pkg/tools"
@@ -131,7 +132,7 @@ func (b *Bus) runWorker() {
 			if !ok {
 				return
 			}
-			resp := b.dispatch(context.Background(), env.req)
+			resp := b.dispatch(env.ctx, env.req)
 			env.reply(resp, nil)
 		case <-b.done:
 			return

pkg/security/securebus/transport.go

@@ -38,6 +38,7 @@ type ChannelTransport struct {
 }
 
 type channelEnvelope struct {
+	ctx    context.Context
 	req    itr.ToolRequest
 	respCh chan channelResult
 }
@@ -70,7 +71,7 @@ func (ct *ChannelTransport) Requests() <-chan channelEnvelope {
 // Send enqueues req and blocks until the response arrives or ctx is cancelled.
 func (ct *ChannelTransport) Send(ctx context.Context, req itr.ToolRequest) (itr.ToolResponse, error) {
 	respCh := make(chan channelResult, 1)
-	env := channelEnvelope{req: req, respCh: respCh}
+	env := channelEnvelope{ctx: ctx, req: req, respCh: respCh}
 
 	select {
 	case ct.reqCh <- env:

pkg/skills/installer.go

@@ -14,7 +14,7 @@ import (
 )
 
 type SkillInstaller struct {
-	workspace string
+	skillsDir string // directory where skills are installed
 }
 
 type AvailableSkill struct {
@@ -31,14 +31,16 @@ type BuiltinSkill struct {
 	Enabled bool   `json:"enabled"`
 }
 
-func NewSkillInstaller(workspace string) *SkillInstaller {
+// NewSkillInstaller creates an installer that manages skills in the given directory.
+// Callers should pass the XDG skills dir (config.SkillsDir()) for new installs.
+func NewSkillInstaller(skillsDir string) *SkillInstaller {
 	return &SkillInstaller{
-		workspace: workspace,
+		skillsDir: skillsDir,
 	}
 }
 
 func (si *SkillInstaller) InstallFromGitHub(ctx context.Context, repo string) error {
-	skillDir := filepath.Join(si.workspace, "skills", filepath.Base(repo))
+	skillDir := filepath.Join(si.skillsDir, filepath.Base(repo))
 
 	if _, err := os.Stat(skillDir); err == nil {
 		return fmt.Errorf("skill '%s' already exists", filepath.Base(repo))
@@ -80,7 +82,7 @@ func (si *SkillInstaller) InstallFromGitHub(ctx context.Context, repo string) er
 }
 
 func (si *SkillInstaller) Uninstall(skillName string) error {
-	skillDir := filepath.Join(si.workspace, "skills", skillName)
+	skillDir := filepath.Join(si.skillsDir, skillName)
 
 	if _, err := os.Stat(skillDir); os.IsNotExist(err) {
 		return fmt.Errorf("skill '%s' not found", skillName)
@@ -126,7 +128,7 @@ func (si *SkillInstaller) ListAvailableSkills(ctx context.Context) ([]AvailableS
 }
 
 func (si *SkillInstaller) ListBuiltinSkills() []BuiltinSkill {
-	builtinSkillsDir := filepath.Join(filepath.Dir(si.workspace), "picoclaw", "skills")
+	builtinSkillsDir := si.skillsDir
 
 	entries, err := os.ReadDir(builtinSkillsDir)
 	if err != nil {

pkg/skills/loader.go

@@ -59,17 +59,19 @@ func (info SkillInfo) validate() error {
 }
 
 type SkillsLoader struct {
-	workspace       string
-	workspaceSkills string // workspace skills (项目级别)
-	globalSkills    string // 全局 skills (~/.picoclaw/skills)
-	builtinSkills   string // 内置 skills
+	primarySkills string // primary skills directory (XDG data dir or workspace/skills)
+	globalSkills  string // user-level override skills (~/.config/picoclaw/skills)
+	builtinSkills string // built-in skills (bundled with binary)
 }
 
-func NewSkillsLoader(workspace string, globalSkills string, builtinSkills string) *SkillsLoader {
+// NewSkillsLoader creates a loader that searches for skills in three directories
+// with priority: primary > global > builtin. The primary directory is typically
+// $XDG_DATA_HOME/picoclaw/skills; the global directory allows user overrides;
+// and the builtin directory ships with the binary.
+func NewSkillsLoader(primarySkillsDir string, globalSkills string, builtinSkills string) *SkillsLoader {
 	return &SkillsLoader{
-		workspace:       workspace,
-		workspaceSkills: filepath.Join(workspace, "skills"),
-		globalSkills:    globalSkills, // ~/.picoclaw/skills
+		primarySkills: primarySkillsDir,
+		globalSkills:  globalSkills,
 		builtinSkills: builtinSkills,
 	}
 }
@@ -77,11 +79,11 @@ func NewSkillsLoader(workspace string, globalSkills string, builtinSkills string
 func (sl *SkillsLoader) ListSkills() []SkillInfo {
 	skills := make([]SkillInfo, 0)
 
-	if sl.workspaceSkills != "" {
-		if dirs, err := os.ReadDir(sl.workspaceSkills); err == nil {
+	if sl.primarySkills != "" {
+		if dirs, err := os.ReadDir(sl.primarySkills); err == nil {
 			for _, dir := range dirs {
 				if dir.IsDir() {
-					skillFile := filepath.Join(sl.workspaceSkills, dir.Name(), "SKILL.md")
+					skillFile := filepath.Join(sl.primarySkills, dir.Name(), "SKILL.md")
 					if _, err := os.Stat(skillFile); err == nil {
 						info := SkillInfo{
 							Name:   dir.Name(),
@@ -193,9 +195,8 @@ func (sl *SkillsLoader) ListSkills() []SkillInfo {
 }
 
 func (sl *SkillsLoader) LoadSkill(name string) (string, bool) {
-	// 1. 优先从 workspace skills 加载（项目级别）
-	if sl.workspaceSkills != "" {
-		skillFile := filepath.Join(sl.workspaceSkills, name, "SKILL.md")
+	if sl.primarySkills != "" {
+		skillFile := filepath.Join(sl.primarySkills, name, "SKILL.md")
 		if content, err := os.ReadFile(skillFile); err == nil {
 			return sl.stripFrontmatter(string(content)), true
 		}

pkg/state/state.go

@@ -88,7 +88,8 @@ func NewManager(workspace string, opts ...Option) *Manager {
 }
 
 func (sm *Manager) loadFromDelegate() {
-	ctx := context.Background()
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
 	if v, err := sm.delegate.GetKV(ctx, kvAgentID, "state:last_channel"); err == nil && v != "" {
 		sm.state.LastChannel = v
 	}
@@ -103,38 +104,37 @@ func (sm *Manager) loadFromDelegate() {
 }
 
 // SetLastChannel atomically updates the last channel and saves the state.
-func (sm *Manager) SetLastChannel(channel string) error {
+func (sm *Manager) SetLastChannel(ctx context.Context, channel string) error {
 	sm.mu.Lock()
 	defer sm.mu.Unlock()
 
 	sm.state.LastChannel = channel
 	sm.state.Timestamp = time.Now()
 
-	return sm.persist()
+	return sm.persist(ctx)
 }
 
 // SetLastChatID atomically updates the last chat ID and saves the state.
-func (sm *Manager) SetLastChatID(chatID string) error {
+func (sm *Manager) SetLastChatID(ctx context.Context, chatID string) error {
 	sm.mu.Lock()
 	defer sm.mu.Unlock()
 
 	sm.state.LastChatID = chatID
 	sm.state.Timestamp = time.Now()
 
-	return sm.persist()
+	return sm.persist(ctx)
 }
 
 // persist writes the current state to the delegate (KV) or file.
 // Must be called with the lock held.
-func (sm *Manager) persist() error {
+func (sm *Manager) persist(ctx context.Context) error {
 	if sm.delegate != nil {
-		return sm.persistToDelegate()
+		return sm.persistToDelegate(ctx)
 	}
 	return sm.saveAtomic()
 }
 
-func (sm *Manager) persistToDelegate() error {
-	ctx := context.Background()
+func (sm *Manager) persistToDelegate(ctx context.Context) error {
 	ts := sm.state.Timestamp.Format(time.RFC3339Nano)
 
 	if err := sm.delegate.UpsertKV(ctx, kvAgentID, "state:last_channel", sm.state.LastChannel); err != nil {

pkg/state/state_test.go

@@ -1,6 +1,7 @@
 package state
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"path/filepath"
@@ -20,7 +21,7 @@ func TestAtomicSave(t *testing.T) {
 	sm := NewManager(tmpDir)
 
 	// Test SetLastChannel
-	err = sm.SetLastChannel("test-channel")
+	err = sm.SetLastChannel(context.Background(), "test-channel")
 	if err != nil {
 		t.Fatalf("SetLastChannel failed: %v", err)
 	}
@@ -59,7 +60,7 @@ func TestSetLastChatID(t *testing.T) {
 	sm := NewManager(tmpDir)
 
 	// Test SetLastChatID
-	err = sm.SetLastChatID("test-chat-id")
+	err = sm.SetLastChatID(context.Background(), "test-chat-id")
 	if err != nil {
 		t.Fatalf("SetLastChatID failed: %v", err)
 	}
@@ -92,7 +93,7 @@ func TestAtomicity_NoCorruptionOnInterrupt(t *testing.T) {
 	sm := NewManager(tmpDir)
 
 	// Write initial state
-	err = sm.SetLastChannel("initial-channel")
+	err = sm.SetLastChannel(context.Background(), "initial-channel")
 	if err != nil {
 		t.Fatalf("SetLastChannel failed: %v", err)
 	}
@@ -114,7 +115,7 @@ func TestAtomicity_NoCorruptionOnInterrupt(t *testing.T) {
 	os.Remove(tempFile)
 
 	// Now do a proper save
-	err = sm.SetLastChannel("new-channel")
+	err = sm.SetLastChannel(context.Background(), "new-channel")
 	if err != nil {
 		t.Fatalf("SetLastChannel failed: %v", err)
 	}
@@ -139,7 +140,7 @@ func TestConcurrentAccess(t *testing.T) {
 	for i := 0; i < 10; i++ {
 		go func(idx int) {
 			channel := fmt.Sprintf("channel-%d", idx)
-			sm.SetLastChannel(channel)
+			sm.SetLastChannel(context.Background(), channel)
 			done <- true
 		}(i)
 	}
@@ -177,8 +178,8 @@ func TestNewManager_ExistingState(t *testing.T) {
 
 	// Create initial state
 	sm1 := NewManager(tmpDir)
-	sm1.SetLastChannel("existing-channel")
-	sm1.SetLastChatID("existing-chat-id")
+	sm1.SetLastChannel(context.Background(), "existing-channel")
+	sm1.SetLastChatID(context.Background(), "existing-chat-id")
 
 	// Create new manager with same workspace
 	sm2 := NewManager(tmpDir)