初始化仓库
This commit is contained in:
+411
@@ -0,0 +1,411 @@
|
||||
package pi
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"github.com/chenhg5/cc-connect/core"
|
||||
)
|
||||
|
||||
func init() {
|
||||
core.RegisterAgent("pi", New)
|
||||
}
|
||||
|
||||
// Agent drives the pi coding agent CLI (`pi --mode json --no-input`).
|
||||
type Agent struct {
|
||||
cmd string // path to pi binary
|
||||
workDir string
|
||||
model string
|
||||
mode string // "default" | "yolo"
|
||||
thinking string // reasoning effort: off, minimal, low, medium, high, xhigh
|
||||
sessionEnv []string
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
func New(opts map[string]any) (core.Agent, error) {
|
||||
workDir, _ := opts["work_dir"].(string)
|
||||
if workDir == "" {
|
||||
workDir = "."
|
||||
}
|
||||
model, _ := opts["model"].(string)
|
||||
mode, _ := opts["mode"].(string)
|
||||
mode = normalizeMode(mode)
|
||||
|
||||
cmd, _ := opts["cmd"].(string)
|
||||
if cmd == "" {
|
||||
cmd = "pi"
|
||||
}
|
||||
|
||||
if _, err := exec.LookPath(cmd); err != nil {
|
||||
return nil, fmt.Errorf("pi: '%s' not found in PATH, install with: npm install -g @mariozechner/pi-coding-agent", cmd)
|
||||
}
|
||||
|
||||
return &Agent{
|
||||
cmd: cmd,
|
||||
workDir: workDir,
|
||||
model: model,
|
||||
mode: mode,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func normalizeMode(raw string) string {
|
||||
switch strings.ToLower(strings.TrimSpace(raw)) {
|
||||
case "yolo", "bypass", "auto-approve":
|
||||
return "yolo"
|
||||
default:
|
||||
return "default"
|
||||
}
|
||||
}
|
||||
|
||||
func (a *Agent) Name() string { return "pi" }
|
||||
func (a *Agent) CLIBinaryName() string { return "pi" }
|
||||
func (a *Agent) CLIDisplayName() string { return "Pi" }
|
||||
|
||||
func (a *Agent) SetModel(model string) {
|
||||
a.mu.Lock()
|
||||
defer a.mu.Unlock()
|
||||
a.model = model
|
||||
slog.Info("pi: model changed", "model", model)
|
||||
}
|
||||
|
||||
func (a *Agent) GetModel() string {
|
||||
a.mu.Lock()
|
||||
defer a.mu.Unlock()
|
||||
return a.model
|
||||
}
|
||||
|
||||
func (a *Agent) AvailableModels(_ context.Context) []core.ModelOption {
|
||||
return nil // Pi uses its own model registry; no static list here.
|
||||
}
|
||||
|
||||
func (a *Agent) SetSessionEnv(env []string) {
|
||||
a.mu.Lock()
|
||||
defer a.mu.Unlock()
|
||||
a.sessionEnv = env
|
||||
}
|
||||
|
||||
func (a *Agent) StartSession(ctx context.Context, sessionID string) (core.AgentSession, error) {
|
||||
a.mu.Lock()
|
||||
mode := a.mode
|
||||
model := a.model
|
||||
thinking := a.thinking
|
||||
extraEnv := append([]string{}, a.sessionEnv...)
|
||||
a.mu.Unlock()
|
||||
return newPiSession(ctx, a.cmd, a.workDir, model, mode, thinking, sessionID, extraEnv)
|
||||
}
|
||||
|
||||
func (a *Agent) ListSessions(_ context.Context) ([]core.AgentSessionInfo, error) {
|
||||
sessDir := piSessionDir(a.workDir)
|
||||
if sessDir == "" {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
entries, err := os.ReadDir(sessDir)
|
||||
if err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
return nil, nil
|
||||
}
|
||||
return nil, fmt.Errorf("pi: read session dir: %w", err)
|
||||
}
|
||||
|
||||
var sessions []core.AgentSessionInfo
|
||||
for _, entry := range entries {
|
||||
name := entry.Name()
|
||||
if entry.IsDir() || !strings.HasSuffix(name, ".jsonl") {
|
||||
continue
|
||||
}
|
||||
|
||||
info, err := entry.Info()
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
sessionID, summary, msgCount := scanPiSession(filepath.Join(sessDir, name))
|
||||
if sessionID == "" {
|
||||
continue
|
||||
}
|
||||
|
||||
sessions = append(sessions, core.AgentSessionInfo{
|
||||
ID: sessionID,
|
||||
Summary: summary,
|
||||
MessageCount: msgCount,
|
||||
ModifiedAt: info.ModTime(),
|
||||
})
|
||||
}
|
||||
|
||||
sort.Slice(sessions, func(i, j int) bool {
|
||||
return sessions[i].ModifiedAt.After(sessions[j].ModifiedAt)
|
||||
})
|
||||
|
||||
return sessions, nil
|
||||
}
|
||||
|
||||
func (a *Agent) DeleteSession(_ context.Context, sessionID string) error {
|
||||
sessDir := piSessionDir(a.workDir)
|
||||
if sessDir == "" {
|
||||
return fmt.Errorf("pi: cannot determine session directory")
|
||||
}
|
||||
|
||||
path := findSessionFile(sessDir, sessionID)
|
||||
if path == "" {
|
||||
return fmt.Errorf("pi: session %q not found", sessionID)
|
||||
}
|
||||
return os.Remove(path)
|
||||
}
|
||||
|
||||
func (a *Agent) Stop() error { return nil }
|
||||
|
||||
// ── ModeSwitcher ─────────────────────────────────────────────
|
||||
|
||||
func (a *Agent) SetMode(mode string) {
|
||||
a.mu.Lock()
|
||||
defer a.mu.Unlock()
|
||||
a.mode = normalizeMode(mode)
|
||||
slog.Info("pi: mode changed", "mode", a.mode)
|
||||
}
|
||||
|
||||
func (a *Agent) GetMode() string {
|
||||
a.mu.Lock()
|
||||
defer a.mu.Unlock()
|
||||
return a.mode
|
||||
}
|
||||
|
||||
func (a *Agent) PermissionModes() []core.PermissionModeInfo {
|
||||
return []core.PermissionModeInfo{
|
||||
{Key: "default", Name: "Default", NameZh: "默认", Desc: "Standard permissions", DescZh: "标准权限模式"},
|
||||
{Key: "yolo", Name: "YOLO", NameZh: "全自动", Desc: "Auto-approve all tool calls", DescZh: "自动批准所有工具调用"},
|
||||
}
|
||||
}
|
||||
|
||||
// ── MemoryFileProvider ───────────────────────────────────────
|
||||
|
||||
func (a *Agent) ProjectMemoryFile() string {
|
||||
absDir, err := filepath.Abs(a.workDir)
|
||||
if err != nil {
|
||||
absDir = a.workDir
|
||||
}
|
||||
return filepath.Join(absDir, "AGENTS.md")
|
||||
}
|
||||
|
||||
func (a *Agent) GlobalMemoryFile() string {
|
||||
homeDir, err := os.UserHomeDir()
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
return filepath.Join(homeDir, ".pi", "AGENTS.md")
|
||||
}
|
||||
|
||||
// ── ReasoningEffortSwitcher ──────────────────────────────────
|
||||
|
||||
func (a *Agent) SetReasoningEffort(effort string) {
|
||||
a.mu.Lock()
|
||||
defer a.mu.Unlock()
|
||||
a.thinking = effort
|
||||
slog.Info("pi: thinking level changed", "level", effort)
|
||||
}
|
||||
|
||||
func (a *Agent) GetReasoningEffort() string {
|
||||
a.mu.Lock()
|
||||
defer a.mu.Unlock()
|
||||
return a.thinking
|
||||
}
|
||||
|
||||
func (a *Agent) AvailableReasoningEfforts() []string {
|
||||
return []string{"off", "minimal", "low", "medium", "high", "xhigh"}
|
||||
}
|
||||
|
||||
// ── GetWorkDir (for /status display) ─────────────────────────
|
||||
|
||||
func (a *Agent) GetWorkDir() string { return a.workDir }
|
||||
|
||||
// ── HistoryProvider ──────────────────────────────────────────
|
||||
|
||||
func (a *Agent) GetSessionHistory(_ context.Context, sessionID string, limit int) ([]core.HistoryEntry, error) {
|
||||
sessDir := piSessionDir(a.workDir)
|
||||
if sessDir == "" {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
sessFile := findSessionFile(sessDir, sessionID)
|
||||
if sessFile == "" {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
return readPiHistory(sessFile, limit)
|
||||
}
|
||||
|
||||
// ── SkillProvider ────────────────────────────────────────────
|
||||
|
||||
func (a *Agent) SkillDirs() []string {
|
||||
absDir, err := filepath.Abs(a.workDir)
|
||||
if err != nil {
|
||||
absDir = a.workDir
|
||||
}
|
||||
dirs := []string{filepath.Join(absDir, ".pi", "skills")}
|
||||
if home, err := os.UserHomeDir(); err == nil {
|
||||
dirs = append(dirs, filepath.Join(home, ".pi", "skills"))
|
||||
}
|
||||
return dirs
|
||||
}
|
||||
|
||||
// ── Session helpers ──────────────────────────────────────────
|
||||
|
||||
// findSessionFile locates the .jsonl file for a given session UUID in sessDir.
|
||||
// Session files are named: <timestamp>_<uuid>.jsonl — this function extracts
|
||||
// the UUID portion and matches exactly to avoid partial-match vulnerabilities.
|
||||
func findSessionFile(sessDir, sessionID string) string {
|
||||
entries, err := os.ReadDir(sessDir)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
for _, entry := range entries {
|
||||
name := entry.Name()
|
||||
if entry.IsDir() || !strings.HasSuffix(name, ".jsonl") {
|
||||
continue
|
||||
}
|
||||
// Extract UUID: strip .jsonl, then take everything after the last "_".
|
||||
base := strings.TrimSuffix(name, ".jsonl")
|
||||
if idx := strings.LastIndex(base, "_"); idx >= 0 {
|
||||
if base[idx+1:] == sessionID {
|
||||
return filepath.Join(sessDir, name)
|
||||
}
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// piSessionDir returns the pi session directory for the given workDir.
|
||||
// Pi encodes the absolute path as: replace "/" with "-", wrap with "--".
|
||||
// e.g. /home/user/project → --home-user-project--
|
||||
func piSessionDir(workDir string) string {
|
||||
homeDir, err := os.UserHomeDir()
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
absDir, err := filepath.Abs(workDir)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
encoded := "--" + strings.ReplaceAll(strings.TrimPrefix(absDir, "/"), "/", "-") + "--"
|
||||
return filepath.Join(homeDir, ".pi", "agent", "sessions", encoded)
|
||||
}
|
||||
|
||||
// scanPiSession reads a pi session .jsonl file and extracts the session ID,
|
||||
// a summary (first user message), and a message count.
|
||||
func scanPiSession(path string) (sessionID, summary string, msgCount int) {
|
||||
f, err := os.Open(path)
|
||||
if err != nil {
|
||||
return "", "", 0
|
||||
}
|
||||
defer f.Close()
|
||||
|
||||
scanner := bufio.NewScanner(f)
|
||||
scanner.Buffer(make([]byte, 0, 64*1024), 1*1024*1024)
|
||||
|
||||
for scanner.Scan() {
|
||||
var entry map[string]any
|
||||
if err := json.Unmarshal(scanner.Bytes(), &entry); err != nil {
|
||||
continue
|
||||
}
|
||||
switch entry["type"] {
|
||||
case "session":
|
||||
if id, ok := entry["id"].(string); ok {
|
||||
sessionID = id
|
||||
}
|
||||
case "message":
|
||||
msg, _ := entry["message"].(map[string]any)
|
||||
if msg == nil {
|
||||
continue
|
||||
}
|
||||
role, _ := msg["role"].(string)
|
||||
if role == "user" || role == "assistant" {
|
||||
msgCount++
|
||||
}
|
||||
// Use first user message as summary.
|
||||
if role == "user" && summary == "" {
|
||||
content, _ := msg["content"].([]any)
|
||||
for _, c := range content {
|
||||
item, _ := c.(map[string]any)
|
||||
if item != nil {
|
||||
if text, ok := item["text"].(string); ok && text != "" {
|
||||
summary = text
|
||||
runes := []rune(summary)
|
||||
if len(runes) > 80 {
|
||||
summary = string(runes[:80]) + "..."
|
||||
}
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if err := scanner.Err(); err != nil {
|
||||
slog.Warn("pi: scan session error", "path", path, "error", err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// readPiHistory reads user/assistant messages from a pi session file.
|
||||
func readPiHistory(path string, limit int) ([]core.HistoryEntry, error) {
|
||||
f, err := os.Open(path)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer f.Close()
|
||||
|
||||
scanner := bufio.NewScanner(f)
|
||||
scanner.Buffer(make([]byte, 0, 64*1024), 1*1024*1024)
|
||||
|
||||
var all []core.HistoryEntry
|
||||
for scanner.Scan() {
|
||||
var entry map[string]any
|
||||
if err := json.Unmarshal(scanner.Bytes(), &entry); err != nil {
|
||||
continue
|
||||
}
|
||||
if entry["type"] != "message" {
|
||||
continue
|
||||
}
|
||||
msg, _ := entry["message"].(map[string]any)
|
||||
if msg == nil {
|
||||
continue
|
||||
}
|
||||
role, _ := msg["role"].(string)
|
||||
if role != "user" && role != "assistant" {
|
||||
continue
|
||||
}
|
||||
|
||||
var text string
|
||||
content, _ := msg["content"].([]any)
|
||||
for _, c := range content {
|
||||
item, _ := c.(map[string]any)
|
||||
if item != nil {
|
||||
if t, ok := item["text"].(string); ok && t != "" {
|
||||
text = t
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
if text == "" {
|
||||
continue
|
||||
}
|
||||
all = append(all, core.HistoryEntry{Role: role, Content: text})
|
||||
}
|
||||
if err := scanner.Err(); err != nil {
|
||||
return nil, fmt.Errorf("pi: read history: %w", err)
|
||||
}
|
||||
|
||||
if limit > 0 && len(all) > limit {
|
||||
all = all[len(all)-limit:]
|
||||
}
|
||||
return all, nil
|
||||
}
|
||||
+1267
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,503 @@
|
||||
package pi
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/chenhg5/cc-connect/core"
|
||||
)
|
||||
|
||||
// piSession manages a multi-turn pi coding agent conversation.
|
||||
// Each Send() spawns `pi --mode json -p <prompt>`.
|
||||
// Subsequent turns use `--session <sessionID>` to resume.
|
||||
type piSession struct {
|
||||
cmd string
|
||||
workDir string
|
||||
model string
|
||||
mode string
|
||||
thinking string // reasoning effort level for --thinking flag
|
||||
extraEnv []string
|
||||
attachDir string
|
||||
events chan core.Event
|
||||
sessionID atomic.Value // stores string
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
wg sync.WaitGroup
|
||||
alive atomic.Bool
|
||||
|
||||
thinkingBuf strings.Builder // accumulates thinking_delta chunks
|
||||
}
|
||||
|
||||
func newPiSession(ctx context.Context, cmd, workDir, model, mode, thinking, resumeID string, extraEnv []string) (*piSession, error) {
|
||||
sessionCtx, cancel := context.WithCancel(ctx)
|
||||
|
||||
s := &piSession{
|
||||
cmd: cmd,
|
||||
workDir: workDir,
|
||||
model: model,
|
||||
mode: mode,
|
||||
thinking: thinking,
|
||||
extraEnv: extraEnv,
|
||||
attachDir: filepath.Join(workDir, ".cc-connect", "attachments",
|
||||
fmt.Sprintf("pi_%d", time.Now().UnixNano())),
|
||||
events: make(chan core.Event, 64),
|
||||
ctx: sessionCtx,
|
||||
cancel: cancel,
|
||||
}
|
||||
s.alive.Store(true)
|
||||
|
||||
if resumeID != "" && resumeID != core.ContinueSession {
|
||||
s.sessionID.Store(resumeID)
|
||||
}
|
||||
|
||||
return s, nil
|
||||
}
|
||||
|
||||
func (s *piSession) Send(prompt string, images []core.ImageAttachment, files []core.FileAttachment) error {
|
||||
// Keep attachments isolated per session so concurrent sessions in the same
|
||||
// workDir cannot delete files that another Pi process still references.
|
||||
cleanAttachments(s.attachDir)
|
||||
|
||||
// Save all attachments to disk — pi reads them via @file syntax.
|
||||
var atFiles []string
|
||||
if len(images) > 0 {
|
||||
atFiles = append(atFiles, saveImagesToDisk(s.attachDir, images)...)
|
||||
}
|
||||
if len(files) > 0 {
|
||||
atFiles = append(atFiles, saveFilesToDisk(s.attachDir, files)...)
|
||||
}
|
||||
if !s.alive.Load() {
|
||||
return fmt.Errorf("session is closed")
|
||||
}
|
||||
|
||||
args := []string{"--mode", "json", "-p"}
|
||||
|
||||
sid := s.CurrentSessionID()
|
||||
if sid != "" {
|
||||
args = append(args, "--session", sid)
|
||||
}
|
||||
|
||||
if s.model != "" {
|
||||
args = append(args, "--model", s.model)
|
||||
}
|
||||
|
||||
if s.mode == "yolo" {
|
||||
args = append(args, "--auto-approve")
|
||||
}
|
||||
|
||||
if s.thinking != "" {
|
||||
args = append(args, "--thinking", s.thinking)
|
||||
}
|
||||
|
||||
// Pass attachments as @file arguments
|
||||
for _, f := range atFiles {
|
||||
args = append(args, "@"+f)
|
||||
}
|
||||
|
||||
// Append prompt as positional arg
|
||||
args = append(args, prompt)
|
||||
|
||||
slog.Debug("piSession: launching", "resume", sid != "", "args", core.RedactArgs(args))
|
||||
|
||||
cmd := exec.CommandContext(s.ctx, s.cmd, args...)
|
||||
cmd.Dir = s.workDir
|
||||
env := os.Environ()
|
||||
if len(s.extraEnv) > 0 {
|
||||
env = core.MergeEnv(env, s.extraEnv)
|
||||
}
|
||||
cmd.Env = env
|
||||
|
||||
stdout, err := cmd.StdoutPipe()
|
||||
if err != nil {
|
||||
return fmt.Errorf("piSession: stdout pipe: %w", err)
|
||||
}
|
||||
|
||||
var stderrBuf bytes.Buffer
|
||||
cmd.Stderr = &stderrBuf
|
||||
|
||||
if err := cmd.Start(); err != nil {
|
||||
return fmt.Errorf("piSession: start: %w", err)
|
||||
}
|
||||
|
||||
s.wg.Add(1)
|
||||
go s.readLoop(cmd, stdout, &stderrBuf)
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *piSession) readLoop(cmd *exec.Cmd, stdout io.ReadCloser, stderrBuf *bytes.Buffer) {
|
||||
defer s.wg.Done()
|
||||
defer func() {
|
||||
if err := cmd.Wait(); err != nil {
|
||||
stderrMsg := strings.TrimSpace(stderrBuf.String())
|
||||
if stderrMsg != "" {
|
||||
slog.Error("piSession: process failed", "error", err, "stderr", truncStr(stderrMsg, 200))
|
||||
evt := core.Event{Type: core.EventError, Error: fmt.Errorf("%s", stderrMsg)}
|
||||
select {
|
||||
case s.events <- evt:
|
||||
case <-s.ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
// Pi's JSON events are small (typically <1KB each). A 10MB Scanner buffer
|
||||
// is more than sufficient — no need for the bufio.Reader approach used by
|
||||
// adapters that may receive very large single-line responses.
|
||||
scanner := bufio.NewScanner(stdout)
|
||||
scanner.Buffer(make([]byte, 0, 64*1024), 10*1024*1024)
|
||||
|
||||
for scanner.Scan() {
|
||||
line := scanner.Text()
|
||||
if line == "" {
|
||||
continue
|
||||
}
|
||||
|
||||
var raw map[string]any
|
||||
if err := json.Unmarshal([]byte(line), &raw); err != nil {
|
||||
slog.Debug("piSession: non-JSON line", "line", truncStr(line, 100))
|
||||
continue
|
||||
}
|
||||
|
||||
s.handleEvent(raw)
|
||||
}
|
||||
|
||||
if err := scanner.Err(); err != nil {
|
||||
slog.Error("piSession: scanner error", "error", err)
|
||||
evt := core.Event{Type: core.EventError, Error: fmt.Errorf("read stdout: %w", err)}
|
||||
select {
|
||||
case s.events <- evt:
|
||||
case <-s.ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// Emit EventResult when the process finishes.
|
||||
sid := s.CurrentSessionID()
|
||||
evt := core.Event{Type: core.EventResult, SessionID: sid, Done: true}
|
||||
select {
|
||||
case s.events <- evt:
|
||||
case <-s.ctx.Done():
|
||||
}
|
||||
}
|
||||
|
||||
// Pi NDJSON event types:
|
||||
//
|
||||
// session — session metadata with id
|
||||
// agent_start/end — agent lifecycle
|
||||
// turn_start/end — turn boundaries
|
||||
// message_start — beginning of user/assistant/toolResult message
|
||||
// message_update — streaming deltas (assistantMessageEvent sub-events)
|
||||
// message_end — complete message
|
||||
func (s *piSession) handleEvent(raw map[string]any) {
|
||||
eventType, _ := raw["type"].(string)
|
||||
|
||||
switch eventType {
|
||||
case "session":
|
||||
if id, ok := raw["id"].(string); ok && id != "" {
|
||||
s.sessionID.Store(id)
|
||||
slog.Debug("piSession: session started", "session_id", id)
|
||||
}
|
||||
|
||||
case "message_update":
|
||||
s.handleMessageUpdate(raw)
|
||||
|
||||
case "message_end":
|
||||
s.handleMessageEnd(raw)
|
||||
|
||||
case "agent_start", "agent_end", "turn_start", "turn_end", "message_start":
|
||||
// Logged for debugging but no action needed.
|
||||
slog.Debug("piSession: lifecycle event", "type", eventType)
|
||||
|
||||
default:
|
||||
slog.Debug("piSession: unhandled event", "type", eventType)
|
||||
}
|
||||
}
|
||||
|
||||
// handleMessageUpdate processes streaming deltas from pi's assistantMessageEvent.
|
||||
func (s *piSession) handleMessageUpdate(raw map[string]any) {
|
||||
ame, _ := raw["assistantMessageEvent"].(map[string]any)
|
||||
if ame == nil {
|
||||
return
|
||||
}
|
||||
|
||||
subType, _ := ame["type"].(string)
|
||||
|
||||
switch subType {
|
||||
case "text_delta":
|
||||
delta, _ := ame["delta"].(string)
|
||||
if delta != "" {
|
||||
evt := core.Event{Type: core.EventText, Content: delta}
|
||||
select {
|
||||
case s.events <- evt:
|
||||
case <-s.ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
case "thinking_delta":
|
||||
delta, _ := ame["delta"].(string)
|
||||
if delta != "" {
|
||||
s.thinkingBuf.WriteString(delta)
|
||||
}
|
||||
|
||||
case "thinking_end":
|
||||
if s.thinkingBuf.Len() > 0 {
|
||||
evt := core.Event{Type: core.EventThinking, Content: s.thinkingBuf.String()}
|
||||
s.thinkingBuf.Reset()
|
||||
select {
|
||||
case s.events <- evt:
|
||||
case <-s.ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
case "toolcall_end":
|
||||
// Extract tool name and input from the accumulated message content.
|
||||
s.emitToolFromMessage(ame)
|
||||
}
|
||||
}
|
||||
|
||||
// emitToolFromMessage extracts tool call info from a toolcall_end event.
|
||||
func (s *piSession) emitToolFromMessage(ame map[string]any) {
|
||||
msg, _ := ame["message"].(map[string]any)
|
||||
if msg == nil {
|
||||
msg, _ = ame["partial"].(map[string]any)
|
||||
}
|
||||
if msg == nil {
|
||||
return
|
||||
}
|
||||
|
||||
content, _ := msg["content"].([]any)
|
||||
idx := int(0)
|
||||
if ci, ok := ame["contentIndex"].(float64); ok {
|
||||
idx = int(ci)
|
||||
}
|
||||
|
||||
if idx >= 0 && idx < len(content) {
|
||||
item, _ := content[idx].(map[string]any)
|
||||
if item != nil {
|
||||
itemType, _ := item["type"].(string)
|
||||
if itemType == "toolCall" {
|
||||
name, _ := item["name"].(string)
|
||||
input := extractToolInput(item)
|
||||
evt := core.Event{Type: core.EventToolUse, ToolName: name, ToolInput: input}
|
||||
select {
|
||||
case s.events <- evt:
|
||||
case <-s.ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// handleMessageEnd processes completed messages — particularly toolResult messages.
|
||||
func (s *piSession) handleMessageEnd(raw map[string]any) {
|
||||
msg, _ := raw["message"].(map[string]any)
|
||||
if msg == nil {
|
||||
return
|
||||
}
|
||||
|
||||
role, _ := msg["role"].(string)
|
||||
|
||||
switch role {
|
||||
case "toolResult":
|
||||
toolName, _ := msg["toolName"].(string)
|
||||
content, _ := msg["content"].([]any)
|
||||
var output string
|
||||
for _, c := range content {
|
||||
if item, ok := c.(map[string]any); ok {
|
||||
if text, ok := item["text"].(string); ok {
|
||||
output = text
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
evt := core.Event{Type: core.EventToolResult, ToolName: toolName, Content: truncStr(output, 500)}
|
||||
select {
|
||||
case s.events <- evt:
|
||||
case <-s.ctx.Done():
|
||||
return
|
||||
}
|
||||
|
||||
case "assistant":
|
||||
// Check for errors
|
||||
if errMsg, _ := msg["errorMessage"].(string); errMsg != "" {
|
||||
evt := core.Event{Type: core.EventError, Error: fmt.Errorf("%s", errMsg)}
|
||||
select {
|
||||
case s.events <- evt:
|
||||
case <-s.ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// extractToolInput pulls a concise summary from a tool call content item.
|
||||
func extractToolInput(item map[string]any) string {
|
||||
args, _ := item["arguments"].(map[string]any)
|
||||
if args == nil {
|
||||
return ""
|
||||
}
|
||||
// Prefer description or command fields.
|
||||
if desc, ok := args["description"].(string); ok && desc != "" {
|
||||
return desc
|
||||
}
|
||||
if cmd, ok := args["command"].(string); ok && cmd != "" {
|
||||
return cmd
|
||||
}
|
||||
if fp, ok := args["file_path"].(string); ok && fp != "" {
|
||||
return fp
|
||||
}
|
||||
if pattern, ok := args["pattern"].(string); ok && pattern != "" {
|
||||
return pattern
|
||||
}
|
||||
if query, ok := args["query"].(string); ok && query != "" {
|
||||
return query
|
||||
}
|
||||
b, _ := json.Marshal(args)
|
||||
return truncStr(string(b), 200)
|
||||
}
|
||||
|
||||
func (s *piSession) RespondPermission(_ string, _ core.PermissionResult) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *piSession) Events() <-chan core.Event {
|
||||
return s.events
|
||||
}
|
||||
|
||||
func (s *piSession) CurrentSessionID() string {
|
||||
v, _ := s.sessionID.Load().(string)
|
||||
return v
|
||||
}
|
||||
|
||||
func (s *piSession) Alive() bool {
|
||||
return s.alive.Load()
|
||||
}
|
||||
|
||||
func (s *piSession) Close() error {
|
||||
s.alive.Store(false)
|
||||
s.cancel()
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
s.wg.Wait()
|
||||
close(done)
|
||||
}()
|
||||
select {
|
||||
case <-done:
|
||||
close(s.events)
|
||||
case <-time.After(8 * time.Second):
|
||||
slog.Warn("piSession: close timed out, abandoning wg.Wait")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// cleanAttachments removes this session's attachment directory to avoid
|
||||
// accumulating files across turns.
|
||||
func cleanAttachments(attachDir string) {
|
||||
if attachDir == "" {
|
||||
return
|
||||
}
|
||||
if err := os.RemoveAll(attachDir); err != nil {
|
||||
slog.Warn("piSession: failed to clean attachments dir", "dir", attachDir, "error", err)
|
||||
}
|
||||
}
|
||||
|
||||
// saveImagesToDisk saves image attachments to attachDir
|
||||
// and returns the list of absolute file paths.
|
||||
//
|
||||
// img.FileName originates from IM upload metadata and is treated as
|
||||
// untrusted: directory components are stripped (both `/` and `\`, the
|
||||
// latter so Linux strips Windows-style paths too) before joining into
|
||||
// attachDir. Without this, FileName="../../escape.png" wrote to
|
||||
// workDir/escape.png — outside the intended attachments directory.
|
||||
func saveImagesToDisk(attachDir string, images []core.ImageAttachment) []string {
|
||||
if err := os.MkdirAll(attachDir, 0o755); err != nil {
|
||||
slog.Error("piSession: failed to create attachments dir", "error", err)
|
||||
return nil
|
||||
}
|
||||
|
||||
var paths []string
|
||||
for i, img := range images {
|
||||
ext := ".png"
|
||||
switch img.MimeType {
|
||||
case "image/jpeg":
|
||||
ext = ".jpg"
|
||||
case "image/gif":
|
||||
ext = ".gif"
|
||||
case "image/webp":
|
||||
ext = ".webp"
|
||||
}
|
||||
fname := sanitizePiAttachmentName(img.FileName)
|
||||
if fname == "" {
|
||||
fname = fmt.Sprintf("image_%d_%d%s", time.Now().UnixMilli(), i, ext)
|
||||
}
|
||||
fpath := filepath.Join(attachDir, fname)
|
||||
if err := os.WriteFile(fpath, img.Data, 0o644); err != nil {
|
||||
slog.Error("piSession: save image failed", "error", err)
|
||||
continue
|
||||
}
|
||||
paths = append(paths, fpath)
|
||||
}
|
||||
return paths
|
||||
}
|
||||
|
||||
func saveFilesToDisk(attachDir string, files []core.FileAttachment) []string {
|
||||
if err := os.MkdirAll(attachDir, 0o755); err != nil {
|
||||
slog.Error("piSession: failed to create attachments dir", "error", err)
|
||||
return nil
|
||||
}
|
||||
|
||||
paths := make([]string, 0, len(files))
|
||||
for i, f := range files {
|
||||
fname := sanitizePiAttachmentName(f.FileName)
|
||||
if fname == "" {
|
||||
fname = fmt.Sprintf("file_%d_%d", time.Now().UnixMilli(), i)
|
||||
}
|
||||
fpath := filepath.Join(attachDir, fname)
|
||||
if err := os.WriteFile(fpath, f.Data, 0o644); err != nil {
|
||||
slog.Error("piSession: save file failed", "error", err)
|
||||
continue
|
||||
}
|
||||
paths = append(paths, fpath)
|
||||
}
|
||||
return paths
|
||||
}
|
||||
|
||||
// sanitizePiAttachmentName reduces a user-supplied attachment filename to a
|
||||
// safe basename for joining into an attachment directory. Strips directory
|
||||
// components (handling both `/` and `\` so an attacker can't bypass via
|
||||
// Windows-style separators on Linux), and rejects parent / current-directory
|
||||
// references so the caller's empty-name fallback can substitute a generated
|
||||
// name. Mirrors core.SaveFilesToDisk's sanitization.
|
||||
func sanitizePiAttachmentName(name string) string {
|
||||
name = strings.ReplaceAll(name, "\\", "/")
|
||||
name = filepath.Base(name)
|
||||
if name == "" || name == "." || name == ".." {
|
||||
return ""
|
||||
}
|
||||
return name
|
||||
}
|
||||
|
||||
func truncStr(s string, maxRunes int) string {
|
||||
if utf8.RuneCountInString(s) <= maxRunes {
|
||||
return s
|
||||
}
|
||||
return string([]rune(s)[:maxRunes]) + "..."
|
||||
}
|
||||
Reference in New Issue
Block a user