magic-agent 2.4.2 → 2.6.0

package/internal/agent/trae_exec.go

@@ -0,0 +1,433 @@
+package agent
+
+// trae_exec.go - TRAE CLI 2.0 的 `exec --json` 适配。
+//
+// 参数（`traex exec --help` / `exec resume --help` 实测 0.211.1）：
+//
+//	exec --json --skip-git-repo-check [-C ws] [-m model] <tools> <prompt>
+//	exec resume <id>|--last --json --skip-git-repo-check [-m model] <tools> <prompt>
+//
+// 事件流与 Codex 同源（二进制里就是 codex_core）：
+//
+//	{"type":"thread.started","thread_id":"…"}                     — 会话 id
+//	{"type":"turn.started"}
+//	{"type":"item.started","item":{"id":"item_0","type":"command_execution","command":"…"}}
+//	{"type":"item.completed","item":{"id":"item_0","type":"command_execution",
+//	   "aggregated_output":"…","exit_code":0}}
+//	{"type":"item.completed","item":{"id":"item_1","type":"agent_message","text":"…"}}
+//	{"type":"item.completed","item":{"id":"item_2","type":"reasoning","text":"…"}}
+//	{"type":"turn.completed","usage":{"input_tokens":…,"output_tokens":…}}
+//	{"type":"turn.failed","error":{"message":"…"}}
+//
+// 没有 --query-timeout：超时由外层 runner 杀进程组。
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"os"
+	"strings"
+	"time"
+)
+
+// traeDefaultModelFromTOML 从 ~/.trae/traecli.toml 取顶层 model。
+// 只认 `model = "..."` 这一种形态（与 `-c model="o3"` 的配置键一致）；
+// `model_provider = "..."` 之类不会被误命中。
+func traeDefaultModelFromTOML(path string) string {
+	data, err := os.ReadFile(path)
+	if err != nil {
+		return ""
+	}
+	for _, ln := range strings.Split(string(data), "\n") {
+		t := strings.TrimSpace(ln)
+		if t == "" || strings.HasPrefix(t, "#") || !strings.HasPrefix(t, "model") {
+			continue
+		}
+		rest := strings.TrimSpace(strings.TrimPrefix(t, "model"))
+		if !strings.HasPrefix(rest, "=") {
+			continue
+		}
+		v := strings.TrimSpace(strings.TrimPrefix(rest, "="))
+		if i := strings.Index(v, "#"); i >= 0 {
+			v = strings.TrimSpace(v[:i])
+		}
+		v = strings.Trim(v, `"'`)
+		if v != "" {
+			return v
+		}
+	}
+	return ""
+}
+
+// buildTraeExecArgs 构造 2.0 的 traex 参数。
+func (e *TraeEngine) buildTraeExecArgs(req Request, prompt string) []string {
+	resuming := req.SessionID != "" || req.Continue
+	var args []string
+	switch {
+	case req.SessionID != "":
+		args = append(args, "exec", "resume", req.SessionID)
+	case req.Continue:
+		args = append(args, "exec", "resume", "--last")
+	default:
+		args = append(args, "exec")
+	}
+	args = append(args, "--json", "--skip-git-repo-check")
+	if m := traeModelFor(req, e); m != "" {
+		args = append(args, "-m", m)
+	}
+	// resume 子命令没有 -C；续接轮靠子进程 cwd（调用点已设置）。
+	if ws := strings.TrimSpace(req.Workspace); ws != "" && !resuming {
+		args = append(args, "-C", ws)
+	}
+	// 工具模式：off = 逐个 --disallowed-tool 减法；on / 白名单 = -y 全放行
+	//（trae 表达不了「只允许这几个」）。
+	tools := toolsOrDefault(req.Tools)
+	if tools.IsOff() {
+		for _, t := range traeOffDisableTools {
+			args = append(args, "--disallowed-tool", t)
+		}
+	} else {
+		args = append(args, "-y")
+	}
+	// MaxTokens / Temperature 仍不支持（与其余 CLI 同一矩阵②）。
+	_ = req.MaxTokens
+	_ = req.Temperature
+	args = append(args, prompt)
+	return args
+}
+
+// completeExec 非流式：跑一次 `exec --json`，把 JSONL 聚合为最终正文。
+func (e *TraeEngine) completeExec(ctx context.Context, bin string, req Request, start time.Time) (Response, error) {
+	prompt, err := traePrompt(req)
+	if err != nil {
+		return Response{}, err
+	}
+	args := e.buildTraeExecArgs(req, prompt)
+	stdout, stderr, runErr := runCLIIn(ctx, req.Workspace, bin, args...)
+
+	acc := &traeExecAcc{}
+	for _, line := range strings.Split(stdout, "\n") {
+		_ = acc.handleLine(line)
+	}
+	if acc.turnErr != nil {
+		return Response{}, acc.turnErr
+	}
+	if runErr != nil {
+		return Response{}, wrapCliError("trae", stdout, stderr, runErr)
+	}
+	text := strings.TrimSpace(acc.Text.String())
+	if text == "" {
+		return Response{}, fmt.Errorf("trae CLI returned empty result")
+	}
+	sessionID := acc.SessionID
+	if sessionID == "" {
+		sessionID = req.SessionID
+	}
+	in, out := acc.tokenCounts()
+	return Response{
+		Engine:       e.Name(),
+		Text:         text,
+		Model:        e.modelLabel(req),
+		SessionID:    sessionID,
+		InputTokens:  in,
+		OutputTokens: out,
+		TotalTokens:  in + out,
+		Latency:      time.Since(start),
+	}, nil
+}
+
+// streamExec 流式：逐行消费 `exec --json`，正文以增量事件转发。
+func (e *TraeEngine) streamExec(ctx context.Context, bin string, req Request, start time.Time, onEvent func(StreamEvent)) (StreamResult, error) {
+	prompt, err := traePrompt(req)
+	if err != nil {
+		return StreamResult{}, err
+	}
+	args := e.buildTraeExecArgs(req, prompt)
+	acc := &traeExecAcc{OnEvent: onEvent, OnSessionID: req.OnSessionID}
+	err = runStreamCLIIn(ctx, req.Workspace, bin, args, acc.handleLine)
+	if acc.turnErr != nil {
+		acc.emitTurnFailed(acc.SessionID, acc.turnErr)
+		return StreamResult{}, acc.turnErr
+	}
+	if err != nil {
+		acc.emitTurnFailed(acc.SessionID, err)
+		return StreamResult{}, err
+	}
+	text := strings.TrimSpace(acc.Text.String())
+	if text == "" {
+		err := fmt.Errorf("trae CLI returned empty result")
+		acc.emitTurnFailed(acc.SessionID, err)
+		return StreamResult{}, err
+	}
+	if !acc.sawTurnEnd {
+		err := fmt.Errorf("trae CLI stream ended without turn.completed")
+		acc.emitTurnFailed(acc.SessionID, err)
+		return StreamResult{}, err
+	}
+	in, out := acc.tokenCounts()
+	return StreamResult{
+		Response: Response{
+			Engine:       e.Name(),
+			Text:         text,
+			Model:        e.modelLabel(req),
+			SessionID:    acc.SessionID,
+			InputTokens:  in,
+			OutputTokens: out,
+			TotalTokens:  in + out,
+			Latency:      time.Since(start),
+		},
+		Thinking: strings.TrimSpace(acc.Thinking.String()),
+	}, nil
+}
+
+// traeExecAcc 把 2.0 的 JSONL 事件聚合为正文 / 思考 / 工具事件。
+//
+// 正文用「前缀差量」消费：item.updated 可能给**累积全文**，item.completed
+// 给最终全文；同一条 item 的文本只要与已见部分前缀一致就只发差量，
+// 对「增量形状」也天然兼容（不构成前缀时按追加处理）。
+type traeExecAcc struct {
+	OnEvent     func(StreamEvent)
+	OnSessionID func(string)
+
+	SessionID string
+	Text      strings.Builder
+	Thinking  strings.Builder
+	Usage     codexTurnUsage
+
+	sawTurnEnd bool
+	turnErr    error
+
+	curTextID    string
+	curTextSeen  string
+	curThinkID   string
+	curThinkSeen string
+}
+
+func (a *traeExecAcc) emit(ev StreamEvent) {
+	if a.OnEvent != nil {
+		a.OnEvent(ev)
+	}
+}
+
+func (a *traeExecAcc) tokenCounts() (int, int) {
+	in := a.Usage.InputTokens + a.Usage.CachedInputTokens
+	out := a.Usage.OutputTokens + a.Usage.ReasoningOutputTokens
+	return in, out
+}
+
+func (a *traeExecAcc) emitTurnEnd(text string) {
+	if a.OnEvent == nil {
+		return
+	}
+	a.OnEvent(StreamEvent{Kind: KindTurnEnd, Text: text, SessionID: a.SessionID})
+}
+
+func (a *traeExecAcc) emitTurnFailed(sessionID string, err error) {
+	if a.OnEvent == nil || err == nil {
+		return
+	}
+	a.OnEvent(StreamEvent{
+		Kind:      KindTurnFailed,
+		SessionID: sessionID,
+		Error:     err.Error(),
+		Reason:    ReasonOf(err),
+		Status:    StatusFailed,
+	})
+}
+
+// handleLine 消费一行 JSONL；turn.failed / error 返回错误以立即中止进程。
+func (a *traeExecAcc) handleLine(line string) error {
+	s := strings.TrimSpace(line)
+	if s == "" {
+		return nil
+	}
+	var head struct {
+		Type     string `json:"type"`
+		ThreadID string `json:"thread_id"`
+	}
+	if json.Unmarshal([]byte(s), &head) != nil {
+		return nil // 非 JSON 行（警告 / 杂讯）跳过
+	}
+	// 任意带 thread_id 的事件都可以定下会话锚点（thread.started 必带）。
+	if head.ThreadID != "" && a.SessionID == "" {
+		a.SessionID = head.ThreadID
+		if a.OnSessionID != nil {
+			a.OnSessionID(head.ThreadID)
+		}
+	}
+	switch head.Type {
+	case "item.started", "item.updated", "item.completed":
+		a.handleItem(s, head.Type)
+	case "turn.completed":
+		var tc codexTurnCompleted
+		if json.Unmarshal([]byte(s), &tc) == nil {
+			a.Usage = tc.Usage
+		}
+		a.sawTurnEnd = true
+		a.emitTurnEnd(strings.TrimSpace(a.Text.String()))
+	case "turn.failed":
+		var tf struct {
+			Error *struct {
+				Message string `json:"message"`
+			} `json:"error"`
+			Message string `json:"message"`
+		}
+		_ = json.Unmarshal([]byte(s), &tf)
+		msg := tf.Message
+		if msg == "" && tf.Error != nil {
+			msg = tf.Error.Message
+		}
+		if msg == "" {
+			msg = "turn failed"
+		}
+		a.sawTurnEnd = true
+		a.turnErr = fmt.Errorf("trae CLI turn failed: %s", msg)
+		return a.turnErr
+	case "error":
+		var ev struct {
+			Message string `json:"message"`
+		}
+		_ = json.Unmarshal([]byte(s), &ev)
+		msg := ev.Message
+		if msg == "" {
+			msg = "unknown error"
+		}
+		a.sawTurnEnd = true
+		a.turnErr = fmt.Errorf("trae CLI error: %s", msg)
+		return a.turnErr
+	}
+	return nil
+}
+
+// handleItem 处理 item.started / item.updated / item.completed。
+//
+//	agent_message      正文
+//	reasoning          思考
+//	command_execution  工具（started → tool_use，completed → tool_result）
+//	mcp_tool_call      MCP 工具（同上）
+//	其余（file_change / todo_list / web_search 等）暂不映射，忽略即可。
+func (a *traeExecAcc) handleItem(raw, evType string) {
+	var ev struct {
+		Item struct {
+			ID               string          `json:"id"`
+			Type             string          `json:"type"`
+			Text             string          `json:"text"`
+			Command          string          `json:"command"`
+			AggregatedOutput string          `json:"aggregated_output"`
+			Tool             string          `json:"tool"`
+			Name             string          `json:"name"`
+			Arguments        json.RawMessage `json:"arguments"`
+			Result           json.RawMessage `json:"result"`
+		} `json:"item"`
+	}
+	if json.Unmarshal([]byte(raw), &ev) != nil {
+		return
+	}
+	it := ev.Item
+	switch it.Type {
+	case "agent_message":
+		a.consumeText(&a.curTextID, &a.curTextSeen, it.ID, it.Text, &a.Text, KindText)
+	case "reasoning":
+		a.consumeText(&a.curThinkID, &a.curThinkSeen, it.ID, it.Text, &a.Thinking, KindThinking)
+	case "command_execution":
+		if evType == "item.started" {
+			a.emit(StreamEvent{Kind: KindToolUse, Name: "command", ID: it.ID, Text: it.Command})
+			return
+		}
+		if evType == "item.completed" {
+			out := it.AggregatedOutput
+			if out == "" {
+				out = it.Command
+			}
+			a.emit(StreamEvent{Kind: KindToolResult, Name: it.ID, ID: it.ID, Text: out})
+		}
+	case "mcp_tool_call":
+		name := it.Tool
+		if name == "" {
+			name = it.Name
+		}
+		if name == "" {
+			name = "mcp"
+		}
+		if evType == "item.started" {
+			a.emit(StreamEvent{Kind: KindToolUse, Name: name, ID: it.ID, Text: string(it.Arguments)})
+			return
+		}
+		if evType == "item.completed" {
+			a.emit(StreamEvent{Kind: KindToolResult, Name: it.ID, ID: it.ID, Text: rawText(it.Result)})
+		}
+	}
+}
+
+// consumeText 消费同一条 item 的文本更新，只把「新增部分」计入累积并发事件。
+func (a *traeExecAcc) consumeText(curID, curSeen *string, id, text string, acc *strings.Builder, kind StreamEventKind) {
+	if text == "" {
+		return
+	}
+	if id == "" {
+		id = *curID
+	}
+	if id == "" {
+		id = "stream"
+	}
+	if id != *curID {
+		*curID = id
+		*curSeen = ""
+	}
+	var delta string
+	switch {
+	case strings.HasPrefix(text, *curSeen):
+		// 累积形状（item.updated / completed 给全文）
+		delta = text[len(*curSeen):]
+		*curSeen = text
+	case strings.HasPrefix(*curSeen, text):
+		// 迟到的旧版本，忽略
+		return
+	default:
+		// 增量形状
+		delta = text
+		*curSeen += text
+	}
+	if delta == "" {
+		return
+	}
+	acc.WriteString(delta)
+	if a.OnEvent != nil {
+		a.OnEvent(StreamEvent{Kind: kind, Text: delta})
+	}
+}
+
+// rawText 从工具结果字段里抽可读文本（字符串 / {content:[{text}]} / {content:"..."}）。
+func rawText(raw json.RawMessage) string {
+	if len(raw) == 0 {
+		return ""
+	}
+	var s string
+	if json.Unmarshal(raw, &s) == nil {
+		return s
+	}
+	var obj struct {
+		Content json.RawMessage `json:"content"`
+	}
+	if json.Unmarshal(raw, &obj) == nil && len(obj.Content) > 0 {
+		if out := extractToolResultText(obj.Content); out != "" {
+			return out
+		}
+	}
+	var blocks []struct {
+		Type string `json:"type"`
+		Text string `json:"text"`
+	}
+	if json.Unmarshal(raw, &blocks) == nil {
+		var sb strings.Builder
+		for _, b := range blocks {
+			if b.Text != "" {
+				sb.WriteString(b.Text)
+			}
+		}
+		if sb.Len() > 0 {
+			return sb.String()
+		}
+	}
+	return string(raw)
+}

package/internal/agent/trae_exec_test.go

@@ -0,0 +1,181 @@
+package agent
+
+// trae_exec_test.go - `exec --json` 参数构造与 JSONL 事件解析的单测。
+//
+// 不落子进程：buildTraeExecArgs / traeExecAcc 都是纯函数式入口，
+// 这样在 Windows 上也能跑（其它引擎的假 CLI 是 sh 脚本，只在 Unix 上可执行）。
+
+import (
+	"os"
+	"path/filepath"
+	"strings"
+	"testing"
+)
+
+func TestTraeBuildExecArgs(t *testing.T) {
+	cases := []struct {
+		name   string
+		req    Request
+		engine *TraeEngine
+		want   []string
+		forbid []string
+	}{
+		{
+			name: "新会话 + 工作目录 + 模型",
+			req: Request{
+				Model:     "MiniMax-M3.1",
+				Workspace: `D:\work`,
+				Messages:  []Message{{Role: "user", Content: "hi"}},
+			},
+			want: []string{"exec", "--json", "--skip-git-repo-check", "-m MiniMax-M3.1", "-C D:\\work", "-y"},
+		},
+		{
+			name: "续接指定会话：exec resume <id>，无 -C",
+			req: Request{
+				SessionID: "th-9",
+				Workspace: `D:\work`,
+				Messages:  []Message{{Role: "user", Content: "hi"}},
+			},
+			want:   []string{"exec resume th-9", "--json", "-y"},
+			forbid: []string{"-C "},
+		},
+		{
+			name: "Continue：exec resume --last",
+			req:  Request{Continue: true, Messages: []Message{{Role: "user", Content: "hi"}}},
+			want: []string{"exec resume --last", "--json", "-y"},
+		},
+		{
+			name: "工具关闭：逐个 --disallowed-tool",
+			req:  Request{Tools: ToolsOff, Messages: []Message{{Role: "user", Content: "hi"}}},
+			want: []string{
+				"exec", "--json", "--disallowed-tool Bash", "--disallowed-tool Read",
+			},
+			forbid: []string{"-y"},
+		},
+		{
+			name:   "引擎注入模型",
+			req:    Request{Messages: []Message{{Role: "user", Content: "hi"}}},
+			engine: &TraeEngine{Model: "trae/inj"},
+			want:   []string{"-m inj"},
+		},
+	}
+	for _, c := range cases {
+		t.Run(c.name, func(t *testing.T) {
+			eng := &TraeEngine{}
+			if c.engine != nil {
+				eng = c.engine
+			}
+			args := eng.buildTraeExecArgs(c.req, "PROMPT")
+			joined := strings.Join(args, " ")
+			if !strings.HasSuffix(joined, " PROMPT") {
+				t.Errorf("prompt 必须在末尾: %q", joined)
+			}
+			for _, w := range c.want {
+				if !strings.Contains(joined, w) {
+					t.Errorf("缺少 %q: %q", w, joined)
+				}
+			}
+			for _, f := range c.forbid {
+				if strings.Contains(joined, f) {
+					t.Errorf("不应出现 %q: %q", f, joined)
+				}
+			}
+			// 1.x / 不支持的参数一个都不能出现。
+			for _, gone := range []string{"--query-timeout", "--output-format", "model.name", "--resume", "-p "} {
+				if strings.Contains(joined, gone) {
+					t.Errorf("不应出现 %q: %q", gone, joined)
+				}
+			}
+		})
+	}
+}
+
+func TestTraeExecAccStream(t *testing.T) {
+	var events []StreamEvent
+	acc := &traeExecAcc{OnEvent: func(ev StreamEvent) { events = append(events, ev) }}
+	lines := []string{
+		`{"type":"thread.started","thread_id":"th-1"}`,
+		`{"type":"turn.started"}`,
+		`{"type":"item.started","item":{"id":"item_0","type":"command_execution","command":"ls"}}`,
+		`{"type":"item.completed","item":{"id":"item_0","type":"command_execution","command":"ls","aggregated_output":"a.txt","exit_code":0}}`,
+		`{"type":"item.updated","item":{"id":"item_1","type":"agent_message","text":"你"}}`,
+		`{"type":"item.updated","item":{"id":"item_1","type":"agent_message","text":"你好"}}`,
+		`{"type":"item.completed","item":{"id":"item_1","type":"agent_message","text":"你好！"}}`,
+		`{"type":"item.completed","item":{"id":"item_2","type":"reasoning","text":"想一下"}}`,
+		`{"type":"turn.completed","usage":{"input_tokens":100,"cached_input_tokens":20,"output_tokens":7,"reasoning_output_tokens":3}}`,
+	}
+	for _, ln := range lines {
+		if err := acc.handleLine(ln); err != nil {
+			t.Fatalf("handleLine(%s): %v", ln, err)
+		}
+	}
+	if got := acc.Text.String(); got != "你好！" {
+		t.Errorf("Text = %q want 你好！", got)
+	}
+	if got := acc.Thinking.String(); got != "想一下" {
+		t.Errorf("Thinking = %q want 想一下", got)
+	}
+	if acc.SessionID != "th-1" {
+		t.Errorf("SessionID = %q want th-1", acc.SessionID)
+	}
+	in, out := acc.tokenCounts()
+	if in != 120 || out != 10 {
+		t.Errorf("tokens = (%d,%d) want (120,10)", in, out)
+	}
+	if !acc.sawTurnEnd {
+		t.Error("sawTurnEnd 应为 true")
+	}
+	var kinds []StreamEventKind
+	for _, ev := range events {
+		kinds = append(kinds, ev.Kind)
+	}
+	want := []StreamEventKind{KindToolUse, KindToolResult, KindText, KindText, KindText, KindThinking, KindTurnEnd}
+	if len(kinds) != len(want) {
+		t.Fatalf("事件序列 = %v want %v", kinds, want)
+	}
+	for i := range want {
+		if kinds[i] != want[i] {
+			t.Errorf("事件[%d] = %v want %v（全序列 %v）", i, kinds[i], want[i], kinds)
+		}
+	}
+	last := events[len(events)-1]
+	if last.Text != "你好！" || last.SessionID != "th-1" {
+		t.Errorf("turn_end = %+v", last)
+	}
+	if events[0].Name != "command" || events[0].Text != "ls" || events[0].ID != "item_0" {
+		t.Errorf("tool_use = %+v", events[0])
+	}
+	if events[1].Text != "a.txt" || events[1].ID != "item_0" {
+		t.Errorf("tool_result = %+v", events[1])
+	}
+}
+
+func TestTraeExecAccTurnFailed(t *testing.T) {
+	acc := &traeExecAcc{}
+	err := acc.handleLine(`{"type":"turn.failed","error":{"message":"quota exceeded"}}`)
+	if err == nil || !strings.Contains(err.Error(), "quota exceeded") {
+		t.Fatalf("turn.failed 应返回错误，实际 %v", err)
+	}
+	if acc.turnErr == nil {
+		t.Error("turnErr 应被记录")
+	}
+}
+
+func TestTraeDefaultModelFromTOML(t *testing.T) {
+	dir := t.TempDir()
+	p := filepath.Join(dir, "traecli.toml")
+	body := "model = \"MiniMax-M3.1-Flash\"\nmodel_provider = \"minimax\"\n"
+	if err := os.WriteFile(p, []byte(body), 0o644); err != nil {
+		t.Fatal(err)
+	}
+	if got := traeDefaultModelFromTOML(p); got != "MiniMax-M3.1-Flash" {
+		t.Errorf("got %q want MiniMax-M3.1-Flash", got)
+	}
+	// model_provider 不能被当成 model
+	if err := os.WriteFile(p, []byte("model_provider = \"x\"\n"), 0o644); err != nil {
+		t.Fatal(err)
+	}
+	if got := traeDefaultModelFromTOML(p); got != "" {
+		t.Errorf("model_provider 误命中: %q", got)
+	}
+}