diff --git a/cmd/picoclaw/internal/onboard/command.go b/cmd/picoclaw/internal/onboard/command.go index ec1012959..e89c15ab7 100644 --- a/cmd/picoclaw/internal/onboard/command.go +++ b/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{ diff --git a/cmd/picoclaw/internal/onboard/helpers.go b/cmd/picoclaw/internal/onboard/helpers.go index 4db8bdc8b..55bbd7de9 100644 --- a/cmd/picoclaw/internal/onboard/helpers.go +++ b/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) } - - // Write file - if err := os.WriteFile(targetPath, data, 0o644); err != nil { - return fmt.Errorf("Failed to write file %s: %w", targetPath, err) + if err := os.WriteFile(targetPath, []byte(content), 0o644); err != nil { + return fmt.Errorf("failed to write file %s: %w", targetPath, err) } + } - return nil - }) - - return err + return nil } diff --git a/pkg/channels/manager.go b/pkg/channels/manager.go index eac7b872a..6037b29b6 100644 --- a/pkg/channels/manager.go +++ b/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 } diff --git a/pkg/channels/manager_test.go b/pkg/channels/manager_test.go index 3b8d28ac3..928823ac9 100644 --- a/pkg/channels/manager_test.go +++ b/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 diff --git a/pkg/tools/shell_process_unix.go b/pkg/tools/shell_process_unix.go index 7b29a81bf..fa96d75da 100644 --- a/pkg/tools/shell_process_unix.go +++ b/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//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) + } +} diff --git a/pkg/tools/shell_timeout_unix_test.go b/pkg/tools/shell_timeout_unix_test.go index 357e1276e..d0d4a5b9b 100644 --- a/pkg/tools/shell_timeout_unix_test.go +++ b/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)