Merge pull request #25 from dj-oyu/codex/add-dismiss-handling-for-task-status

Use composite channel:chat:task keys for task status and dismiss drafts on final message

cmd/picoclaw/internal/onboard/command.go

@@ -1,14 +1,6 @@
 package onboard
 
-import (
-	"embed"
-
-	"github.com/spf13/cobra"
-)
-
-//go:generate cp -r ../../../../workspace .
-//go:embed workspace
-var embeddedFiles embed.FS
+import "github.com/spf13/cobra"
 
 func NewOnboardCommand() *cobra.Command {
 	cmd := &cobra.Command{

cmd/picoclaw/internal/onboard/helpers.go

@@ -2,7 +2,6 @@ package onboard
 
 import (
 	"fmt"
-	"io/fs"
 	"os"
 	"path/filepath"
 
@@ -10,6 +9,30 @@ import (
 	"github.com/sipeed/picoclaw/pkg/config"
 )
 
+var workspaceTemplates = map[string]string{
+	"AGENTS.md": `# Agent Instructions
+
+You are a helpful AI assistant. Be concise, accurate, and friendly.
+`,
+	"IDENTITY.md": `# Identity
+
+## Name
+PicoClaw 🦞
+`,
+	"SOUL.md": `# Soul
+
+I am picoclaw, a lightweight AI assistant powered by AI.
+`,
+	"USER.md": `# User
+
+Information about user goes here.
+`,
+	"memory/MEMORY.md": `# Long-term Memory
+
+This file stores important information that should persist across sessions.
+`,
+}
+
 func onboard() {
 	configPath := internal.GetConfigPath()
 
@@ -54,48 +77,19 @@ func createWorkspaceTemplates(workspace string) {
 }
 
 func copyEmbeddedToTarget(targetDir string) error {
-	// Ensure target directory exists
 	if err := os.MkdirAll(targetDir, 0o755); err != nil {
-		return fmt.Errorf("Failed to create target directory: %w", err)
+		return fmt.Errorf("failed to create target directory: %w", err)
 	}
 
-	// Walk through all files in embed.FS
-	err := fs.WalkDir(embeddedFiles, "workspace", func(path string, d fs.DirEntry, err error) error {
-		if err != nil {
-			return err
-		}
-
-		// Skip directories
-		if d.IsDir() {
-			return nil
-		}
-
-		// Read embedded file
-		data, err := embeddedFiles.ReadFile(path)
-		if err != nil {
-			return fmt.Errorf("Failed to read embedded file %s: %w", path, err)
-		}
-
-		new_path, err := filepath.Rel("workspace", path)
-		if err != nil {
-			return fmt.Errorf("Failed to get relative path for %s: %v\n", path, err)
-		}
-
-		// Build target file path
-		targetPath := filepath.Join(targetDir, new_path)
-
-		// Ensure target file's directory exists
+	for relPath, content := range workspaceTemplates {
+		targetPath := filepath.Join(targetDir, relPath)
 		if err := os.MkdirAll(filepath.Dir(targetPath), 0o755); err != nil {
-			return fmt.Errorf("Failed to create directory %s: %w", filepath.Dir(targetPath), err)
+			return fmt.Errorf("failed to create directory %s: %w", filepath.Dir(targetPath), err)
+		}
+		if err := os.WriteFile(targetPath, []byte(content), 0o644); err != nil {
+			return fmt.Errorf("failed to write file %s: %w", targetPath, err)
 		}
-
-		// Write file
-		if err := os.WriteFile(targetPath, data, 0o644); err != nil {
-			return fmt.Errorf("Failed to write file %s: %w", targetPath, err)
 	}
 
 	return nil
-	})
-
-	return err
 }

pkg/channels/manager.go

@@ -12,6 +12,7 @@ import (
 	"fmt"
 	"hash/fnv"
 	"math"
+	"strings"
 	"sync"
 	"time"
 
@@ -100,7 +101,7 @@ type Manager struct {
 	typingStops     sync.Map // "channel:chatID" → typingEntry
 	reactionUndos   sync.Map // "channel:chatID" → reactionEntry
 	statusMsgIDs    sync.Map // "channel:chatID" → statusMsgEntry (streaming preview)
-	taskMsgIDs      sync.Map // taskID → statusMsgEntry (background task status)
+	taskMsgIDs      sync.Map // "channel:chatID:taskID" → statusMsgEntry (background task status)
 	statusEditTimes sync.Map // key → time.Time — last EditMessage time for throttling
 }
 
@@ -594,6 +595,16 @@ func (m *Manager) handleStatusSend(ctx context.Context, name string, w *channelW
 	// 4. Channel doesn't support SendWithID or editing — drop silently
 }
 
+func taskStatusKey(channel, chatID, taskID string) string {
+	if taskID == "" {
+		return ""
+	}
+	if channel == "" || chatID == "" {
+		return taskID
+	}
+	return channel + ":" + chatID + ":" + taskID
+}
+
 // handleTaskStatusSend processes IsTaskStatus messages (background task status).
 // It reuses a previously tracked task message, or sends a new one via SendWithID.
 // For channels implementing DraftSender, sendMessageDraft is used to avoid "(edited)".
@@ -603,14 +614,27 @@ func (m *Manager) handleTaskStatusSend(ctx context.Context, name string, w *chan
 		return
 	}
 
-	taskKey := msg.TaskID
+	taskKey := taskStatusKey(name, msg.ChatID, msg.TaskID)
 
 	// Final message: send as permanent (non-draft) message so it persists.
 	// Drafts are ephemeral and disappear after a short time; the completion
 	// message must survive. Clear the draft tracking and send via SendWithID
 	// or regular Send, which creates a permanent Telegram message.
 	if msg.Final {
-		m.taskMsgIDs.Delete(taskKey)
+		if v, loaded := m.taskMsgIDs.LoadAndDelete(taskKey); loaded {
+			if entry, ok := v.(statusMsgEntry); ok && entry.draftID != 0 {
+				if drafter, ok := w.ch.(DraftSender); ok {
+					if err := drafter.SendDraft(ctx, msg.ChatID, entry.draftID, ""); err != nil {
+						logger.WarnCF("channels", "Failed to dismiss task draft before final message", map[string]any{
+							"task_id":  taskKey,
+							"chat_id":  msg.ChatID,
+							"draft_id": entry.draftID,
+							"error":    err.Error(),
+						})
+					}
+				}
+			}
+		}
 		m.statusEditTimes.Delete(taskKey)
 		if sender, ok := w.ch.(MessageSenderWithID); ok {
 			if msgID, err := sender.SendWithID(ctx, msg.ChatID, msg.Content); err == nil && msgID != "" {
@@ -992,7 +1016,7 @@ func (m *Manager) runTTLJanitor(ctx context.Context) {
 }
 
 // PromoteStatusToTask moves the tracked streaming status message for the given
-// channel:chatID key into the task message map under taskID. This allows the
+// channel:chatID key into the task message map under channel:chatID:taskID. This allows the
 // next IsTaskStatus publish to edit the streaming bubble instead of creating a
 // new message. Returns true if a status message was found and promoted.
 func (m *Manager) PromoteStatusToTask(statusKey, taskID string) bool {
@@ -1000,6 +1024,13 @@ func (m *Manager) PromoteStatusToTask(statusKey, taskID string) bool {
 	if !loaded {
 		return false
 	}
+
+	parts := strings.SplitN(statusKey, ":", 2)
+	if len(parts) == 2 {
+		m.taskMsgIDs.Store(taskStatusKey(parts[0], parts[1], taskID), v)
+		return true
+	}
+
 	m.taskMsgIDs.Store(taskID, v)
 	return true
 }

pkg/channels/manager_test.go

@@ -1019,7 +1019,10 @@ func TestHandleTaskStatusSend_EditsExisting(t *testing.T) {
 	w := &channelWorker{ch: ch, limiter: rate.NewLimiter(rate.Inf, 1)}
 
 	// Pre-store task message
-	m.taskMsgIDs.Store("task-abc", statusMsgEntry{messageID: "task-msg-1", createdAt: time.Now()})
+	m.taskMsgIDs.Store(
+		taskStatusKey("test", "123", "task-abc"),
+		statusMsgEntry{messageID: "task-msg-1", createdAt: time.Now()},
+	)
 
 	msg := bus.OutboundMessage{
 		Channel:      "test",
@@ -1065,7 +1068,7 @@ func TestHandleTaskStatusSend_SendsNewAndTracks(t *testing.T) {
 		t.Fatal("expected SendWithID to be called")
 	}
 
-	v, ok := m.taskMsgIDs.Load("task-xyz")
+	v, ok := m.taskMsgIDs.Load(taskStatusKey("test", "123", "task-xyz"))
 	if !ok {
 		t.Fatal("expected taskMsgIDs to contain tracked entry")
 	}
@@ -1436,6 +1439,137 @@ func TestHandleTaskStatusSend_UsesDraftSender(t *testing.T) {
 	}
 }
 
+func TestHandleTaskStatusSend_Final_DismissesDraftBeforePermanentMessage(t *testing.T) {
+	m := newTestManager()
+	var dismissCalled bool
+	var dismissDraftID int
+	var dismissContent string
+	var finalSendWithIDCalled bool
+
+	ch := &mockDraftSender{
+		mockChannel: mockChannel{
+			sendFn: func(_ context.Context, _ bus.OutboundMessage) error {
+				t.Fatal("Send should not be called when SendWithID succeeds")
+				return nil
+			},
+		},
+		draftFn: func(_ context.Context, chatID string, draftID int, content string) error {
+			dismissCalled = true
+			dismissDraftID = draftID
+			dismissContent = content
+			if chatID != "123" {
+				t.Fatalf("expected dismiss chatID 123, got %s", chatID)
+			}
+			return nil
+		},
+		editFn: func(_ context.Context, _, _, _ string) error { return nil },
+		sendWithID: func(_ context.Context, chatID, content string) (string, error) {
+			finalSendWithIDCalled = true
+			if chatID != "123" {
+				t.Fatalf("expected final chatID 123, got %s", chatID)
+			}
+			if content != "task completed" {
+				t.Fatalf("expected final content 'task completed', got %s", content)
+			}
+			return "task-final-1", nil
+		},
+	}
+
+	w := &channelWorker{ch: ch, limiter: rate.NewLimiter(rate.Inf, 1)}
+
+	m.taskMsgIDs.Store(taskStatusKey("test", "123", "task-final"), statusMsgEntry{draftID: 42, createdAt: time.Now()})
+	m.statusEditTimes.Store(taskStatusKey("test", "123", "task-final"), time.Now())
+
+	msg := bus.OutboundMessage{
+		Channel:      "test",
+		ChatID:       "123",
+		Content:      "task completed",
+		IsTaskStatus: true,
+		TaskID:       "task-final",
+		Final:        true,
+	}
+	m.handleTaskStatusSend(context.Background(), "test", w, msg)
+
+	if !dismissCalled {
+		t.Fatal("expected SendDraft dismiss call for final task status")
+	}
+	if dismissDraftID != 42 {
+		t.Fatalf("expected dismiss draftID 42, got %d", dismissDraftID)
+	}
+	if dismissContent != "" {
+		t.Fatalf("expected empty dismiss content, got %q", dismissContent)
+	}
+	if !finalSendWithIDCalled {
+		t.Fatal("expected final SendWithID to be called")
+	}
+	if _, loaded := m.taskMsgIDs.Load(taskStatusKey("test", "123", "task-final")); loaded {
+		t.Fatal("expected taskMsgIDs entry to be deleted for final task status")
+	}
+	if _, loaded := m.statusEditTimes.Load(taskStatusKey("test", "123", "task-final")); loaded {
+		t.Fatal("expected statusEditTimes entry to be deleted for final task status")
+	}
+}
+
+func TestHandleTaskStatusSend_DraftStreaming_IsolatedByChatThread(t *testing.T) {
+	m := newTestManager()
+
+	type draftCall struct {
+		chatID  string
+		draftID int
+		content string
+	}
+	calls := make([]draftCall, 0, 2)
+
+	ch := &mockDraftSender{
+		mockChannel: mockChannel{
+			sendFn: func(_ context.Context, _ bus.OutboundMessage) error { return nil },
+		},
+		draftFn: func(_ context.Context, chatID string, draftID int, content string) error {
+			calls = append(calls, draftCall{chatID: chatID, draftID: draftID, content: content})
+			return nil
+		},
+		editFn:     func(_ context.Context, _, _, _ string) error { return nil },
+		sendWithID: func(_ context.Context, _, _ string) (string, error) { return "", nil },
+	}
+
+	w := &channelWorker{ch: ch, limiter: rate.NewLimiter(rate.Inf, 1)}
+
+	msgA := bus.OutboundMessage{
+		Channel:      "test",
+		ChatID:       "-100/10",
+		Content:      "A:10%",
+		IsTaskStatus: true,
+		TaskID:       "shared-task",
+	}
+	msgB := bus.OutboundMessage{
+		Channel:      "test",
+		ChatID:       "-100/20",
+		Content:      "B:10%",
+		IsTaskStatus: true,
+		TaskID:       "shared-task",
+	}
+
+	m.handleTaskStatusSend(context.Background(), "test", w, msgA)
+	m.handleTaskStatusSend(context.Background(), "test", w, msgB)
+
+	if len(calls) != 2 {
+		t.Fatalf("expected 2 SendDraft calls, got %d", len(calls))
+	}
+	if calls[0].chatID == calls[1].chatID {
+		t.Fatalf("expected different chat threads, got %q and %q", calls[0].chatID, calls[1].chatID)
+	}
+	if calls[0].draftID == calls[1].draftID {
+		t.Fatalf("expected distinct draft IDs per thread key, both got %d", calls[0].draftID)
+	}
+
+	if _, loaded := m.taskMsgIDs.Load(taskStatusKey("test", "-100/10", "shared-task")); !loaded {
+		t.Fatal("expected taskMsgIDs entry for thread A")
+	}
+	if _, loaded := m.taskMsgIDs.Load(taskStatusKey("test", "-100/20", "shared-task")); !loaded {
+		t.Fatal("expected taskMsgIDs entry for thread B")
+	}
+}
+
 func TestHandleTaskStatusSend_DraftFailure_DoesNotClobberTrackedMessageID(t *testing.T) {
 	m := newTestManager()
 	var sendWithIDCount int

pkg/tools/shell_process_unix.go

@@ -3,7 +3,10 @@
 package tools
 
 import (
+	"os"
 	"os/exec"
+	"strconv"
+	"strings"
 	"syscall"
 )
 
@@ -26,7 +29,57 @@ func terminateProcessTree(cmd *exec.Cmd) error {
 
 	// Kill the entire process group spawned by the shell command.
 	_ = syscall.Kill(-pid, syscall.SIGKILL)
+	// Some shells/background jobs may still leave descendants around
+	// briefly; aggressively walk /proc and kill child processes too.
+	killDescendants(pid)
 	// Fallback kill on the shell process itself.
 	_ = cmd.Process.Kill()
 	return nil
 }
+
+func killDescendants(ppid int) {
+	if ppid <= 0 {
+		return
+	}
+
+	entries, err := os.ReadDir("/proc")
+	if err != nil {
+		return
+	}
+
+	for _, e := range entries {
+		if !e.IsDir() {
+			continue
+		}
+		childPID, err := strconv.Atoi(e.Name())
+		if err != nil || childPID <= 0 || childPID == ppid {
+			continue
+		}
+
+		statPath := "/proc/" + e.Name() + "/stat"
+		data, err := os.ReadFile(statPath)
+		if err != nil {
+			continue
+		}
+
+		// /proc/<pid>/stat: pid (comm) state ppid ...
+		raw := string(data)
+		end := strings.LastIndex(raw, ")")
+		if end == -1 || end+2 >= len(raw) {
+			continue
+		}
+		fields := strings.Fields(raw[end+2:])
+		if len(fields) < 2 {
+			continue
+		}
+		parent, err := strconv.Atoi(fields[1])
+		if err != nil || parent != ppid {
+			continue
+		}
+
+		// Recurse first, then kill child process/group.
+		killDescendants(childPID)
+		_ = syscall.Kill(-childPID, syscall.SIGKILL)
+		_ = syscall.Kill(childPID, syscall.SIGKILL)
+	}
+}

pkg/tools/shell_timeout_unix_test.go

@@ -13,12 +13,33 @@ import (
 	"time"
 )
 
-func processExists(pid int) bool {
+func processRunning(pid int) bool {
 	if pid <= 0 {
 		return false
 	}
+	// kill(0) can return success for zombie processes too, so inspect /proc
+	// state and treat zombies as not-running for timeout cleanup assertions.
 	err := syscall.Kill(pid, 0)
-	return err == nil || err == syscall.EPERM
+	if err != nil && err != syscall.EPERM {
+		return false
+	}
+
+	data, readErr := os.ReadFile("/proc/" + strconv.Itoa(pid) + "/stat")
+	if readErr != nil {
+		return false
+	}
+	raw := string(data)
+	end := strings.LastIndex(raw, ")")
+	if end == -1 || end+2 >= len(raw) {
+		return true // best effort fallback
+	}
+	fields := strings.Fields(raw[end+2:])
+	if len(fields) == 0 {
+		return true // best effort fallback
+	}
+
+	state := fields[0]
+	return state != "Z"
 }
 
 func TestShellTool_TimeoutKillsChildProcess(t *testing.T) {
@@ -55,7 +76,7 @@ func TestShellTool_TimeoutKillsChildProcess(t *testing.T) {
 
 	deadline := time.Now().Add(2 * time.Second)
 	for time.Now().Before(deadline) {
-		if !processExists(childPID) {
+		if !processRunning(childPID) {
 			return
 		}
 		time.Sleep(50 * time.Millisecond)