diff --git a/.iflow/settings.json b/.iflow/settings.json new file mode 100644 index 000000000..0165dda42 --- /dev/null +++ b/.iflow/settings.json @@ -0,0 +1,86 @@ +{ + "hooks": { + "Notification": [ + { + "hooks": [ + { + "type": "command", + "command": "entire hooks iflow notification" + } + ] + } + ], + "PostToolUse": [ + { + "matcher": "*", + "hooks": [ + { + "type": "command", + "command": "entire hooks iflow post-tool-use" + } + ] + } + ], + "PreToolUse": [ + { + "matcher": "*", + "hooks": [ + { + "type": "command", + "command": "entire hooks iflow pre-tool-use" + } + ] + } + ], + "SessionEnd": [ + { + "type": "command", + "command": "entire hooks iflow session-end" + } + ], + "SessionStart": [ + { + "matcher": "startup", + "hooks": [ + { + "type": "command", + "command": "entire hooks iflow session-start" + } + ] + } + ], + "SetUpEnvironment": [ + { + "type": "command", + "command": "entire hooks iflow set-up-environment" + } + ], + "Stop": [ + { + "type": "command", + "command": "entire hooks iflow stop" + } + ], + "SubagentStop": [ + { + "type": "command", + "command": "entire hooks iflow subagent-stop" + } + ], + "UserPromptSubmit": [ + { + "hooks": [ + { + "type": "command", + "command": "entire hooks iflow user-prompt-submit" + } + ] + } + ] + }, + "permissions": { + "deny": [ + "Read(./.entire/metadata/**)" + ] + } +} diff --git a/IFLOW.md b/IFLOW.md new file mode 100644 index 000000000..675e22f49 --- /dev/null +++ b/IFLOW.md @@ -0,0 +1,244 @@ +# Entire CLI 架构概述 + +## 1. 项目概览 + +Entire CLI 是一个与 Git 工作流集成的工具，用于捕获和管理 AI 代理（Claude Code、Gemini CLI、Cursor、OpenCode 等）的会话数据。核心功能： + +- 记录 AI 代理的完整交互过程（提示词、响应、修改的文件） +- 提供会话恢复和回滚功能 +- 在单独的分支上保存会话元数据，保持主分支 Git 历史整洁 +- 支持审计和合规要求 + +**技术栈**: Go 1.26.x, Cobra CLI, go-git/v6, mise 构建工具 + +## 2. 构建与命令 + +### 开发命令 +```bash +mise install # 安装依赖 +mise run build # 构建 CLI +mise run fmt # 代码格式化 +mise run lint # 代码检查 +``` + +### 测试命令 +```bash +mise run test # 单元测试 +mise run test:integration # 集成测试 +mise run test:ci # CI 完整测试（单元+集成+E2E canary） +mise run test:e2e:canary # E2E canary 测试（Vogon 代理，无 API 调用） +``` + +### 提交前必须执行 +```bash +mise run fmt && mise run lint && mise run test:ci +``` + +## 3. 核心架构 + +### Agent 系统 (`cmd/entire/cli/agent/`) + +Agent 系统使用**控制反转**模式：代理是被动数据提供者，将原生 hook 载荷转换为标准化生命周期事件，框架负责所有编排。 + +**核心接口** (`agent.Agent`): +- `Name()` / `Type()` / `Description()` - 身份标识 +- `DetectPresence()` - 检测代理是否配置 +- `ProtectedDirs()` - 回滚时保护的目录 +- `ReadTranscript()` / `ChunkTranscript()` / `ReassembleTranscript()` - 会话记录处理 +- `ParseHookEvent()` - **核心贡献点** - 将原生 hook 转换为标准化事件 + +**可选接口**: +- `HookSupport` - 自动安装 hooks 到代理配置文件 +- `TranscriptAnalyzer` - 从会话记录提取文件列表和提示词 +- `TokenCalculator` - 计算 token 使用量 +- `SubagentAwareExtractor` - 处理子代理贡献 + +**已支持的代理**: Claude Code, Gemini CLI, Cursor, OpenCode, Factory AI Droid, Vogon (测试用) + +### Strategy 系统 (`cmd/entire/cli/strategy/`) + +Manual-Commit 策略管理会话数据和检查点： +- 创建影子分支 `entire/-` 存储会话元数据 +- 用户提交时压缩到 `entire/checkpoints/v1` 分支 +- 支持 Git worktrees，每个 worktree 独立跟踪 +- 不修改活跃分支历史，安全用于 main/master + +**关键文件**: +- `strategy.go` - 接口定义和上下文结构 +- `manual_commit.go` - 主实现 +- `manual_commit_session.go` - 会话状态管理 +- `manual_commit_condensation.go` - 压缩逻辑 +- `manual_commit_rewind.go` - 回滚实现 + +### Checkpoint 系统 (`cmd/entire/cli/checkpoint/`) + +- `temporary.go` - 影子分支操作 +- `committed.go` - 元数据分支操作 +- 使用 12 位十六进制检查点 ID 进行双向链接 + +### Session 状态机 (`cmd/entire/cli/session/`) + +**阶段**: `ACTIVE` → `IDLE` → `ENDED` + +**事件**: `TurnStart`, `TurnEnd`, `GitCommit`, `SessionStart`, `SessionEnd` + +## 4. 代码风格 + +### Go 代码规范 +- 遵循 `golangci-lint` 检查规则 +- 使用 `slog` 进行结构化日志记录 +- 所有测试使用 `t.Parallel()` 并行执行（除非修改进程全局状态） + +### 错误处理模式 +- `root.go` 设置 `SilenceErrors: true`，Cobra 不打印错误 +- `main.go` 打印错误到 stderr，除非是 `SilentError` +- 使用 `NewSilentError()` 当需要自定义用户友好错误消息时 + +### Settings 访问 +```go +import "github.com/entireio/cli/cmd/entire/cli/settings" + +s, err := settings.Load(ctx) +if s.Enabled { ... } +``` +**禁止**: 直接读取 `.entire/settings.json` + +### 日志 vs 用户输出 +- **内部日志**: `logging.Debug/Info/Warn/Error(ctx, msg, attrs...)` → `.entire/logs/` +- **用户输出**: `fmt.Fprint*(cmd.OutOrStdout(), ...)` +- **隐私**: 不记录用户内容（提示词、文件内容、提交消息） + +### Git 操作注意事项 + +**禁止使用 go-git v5/v6 进行 `checkout` 或 `reset --hard`**： +go-git 有 bug 会错误删除 `.gitignore` 中列出的未跟踪目录。使用 git CLI： +```go +// 正确方式 +cmd := exec.CommandContext(ctx, "git", "reset", "--hard", hash.String()) +``` + +**路径处理**： +始终使用仓库根目录而非 `os.Getwd()`： +```go +// 错误 - 从子目录运行时会出错 +cwd, _ := os.Getwd() +absPath := filepath.Join(cwd, file) + +// 正确 +repoRoot, _ := paths.WorktreeRoot() +absPath := filepath.Join(repoRoot, file) +``` + +## 5. 测试 + +### 测试隔离 +**测试必须使用隔离的临时仓库**，禁止使用真实仓库 CWD： +```go +tmpDir := t.TempDir() +testutil.InitRepo(t, tmpDir) // git init + user config + disable GPG +testutil.WriteFile(t, tmpDir, "f.txt", "init") +testutil.GitAdd(t, tmpDir, "f.txt") +testutil.GitCommit(t, tmpDir, "init") +t.Chdir(tmpDir) +``` + +### 测试文件位置 +- 单元测试: 与源文件同目录 `*_test.go` +- 集成测试: `cmd/entire/cli/integration_test/` (使用 `//go:build integration` 标签) +- E2E 测试: `e2e/tests/` (使用 `//go:build e2e` 标签) + +### 代码重复检查 +```bash +mise run dup # 综合检查（阈值 50） +mise run dup:staged # 仅检查暂存文件 +``` + +## 6. 安全 + +### 数据保护 +- `settings.local.json` 默认不提交到 Git +- 会话元数据存储在单独分支，与主分支分离 +- `redact/` 包处理敏感信息 + +### 权限控制 +- 设置文件使用权限 `0o644` +- 遵循最小权限原则 + +### 遥测 +- 支持匿名使用统计，可通过配置关闭 +- 非阻塞方式收集，不影响核心功能 + +## 7. 配置 + +### 配置文件 +位于 `.entire/` 目录： + +**settings.json** (项目设置，提交到 Git): +```json +{ + "enabled": true, + "commit_linking": "prompt", + "strategy_options": { + "push_sessions": true, + "summarize": { "enabled": true } + } +} +``` + +**settings.local.json** (本地设置，不提交): +```json +{ + "enabled": false, + "log_level": "debug" +} +``` + +### 主要配置选项 + +| 选项 | 值 | 描述 | +|------|-----|------| +| `enabled` | `true`, `false` | 启用/禁用 Entire | +| `log_level` | `debug`, `info`, `warn`, `error` | 日志级别 | +| `commit_linking` | `always`, `prompt` | 提交链接模式 | +| `strategy_options.push_sessions` | `true`, `false` | Git push 时自动推送 checkpoints 分支 | +| `strategy_options.summarize.enabled` | `true`, `false` | 自动生成 AI 摘要 | +| `telemetry` | `true`, `false` | 发送匿名使用统计 | +| `external_agents` | `true`, `false` | 启用外部代理插件发现 | + +## 8. 无障碍支持 + +环境变量 `ACCESSIBLE=1` 启用无障碍模式，使用简单文本提示替代交互式 TUI 元素。 + +在 `cli` 包中使用 `NewAccessibleForm()`，在 `strategy` 包中使用 `isAccessibleMode()` 辅助函数。 + +## 9. 添加新代理 + +参考 `docs/architecture/agent-guide.md` 和 `docs/architecture/agent-integration-checklist.md`。 + +关键步骤： +1. 创建 `cmd/entire/cli/agent/youragent/` 目录 +2. 实现 `agent.Agent` 接口（19 个方法） +3. 实现 `ParseHookEvent()` - 核心贡献点 +4. 根据需要实现可选接口 +5. 通过 `init()` 注册代理 +6. 添加编译时接口断言 +7. 编写测试 + +## 10. 行为规则 + +### 开发流程 +- 非平凡任务（3+ 步骤或架构决策）先进入计划模式 +- 使用子代理保持主上下文窗口清洁 +- 任务完成后验证正确性（运行测试、检查日志） +- 追求优雅但不过度工程化 + +### 任务管理 +1. 先写计划到 `tasks/todo.md` +2. 开始实现前确认计划 +3. 进度跟踪：完成一项标记一项 +4. 记录结果和经验教训 + +### 核心原则 +- **简洁优先**: 每个改动尽可能简单，影响最小代码 +- **不偷懒**: 找到根本原因，不做临时修复 +- **最小影响**: 只触及必要的代码 diff --git a/cmd/entire/cli/agent/iflow/hooks.go b/cmd/entire/cli/agent/iflow/hooks.go new file mode 100644 index 000000000..aa75d0d3f --- /dev/null +++ b/cmd/entire/cli/agent/iflow/hooks.go @@ -0,0 +1,518 @@ +package iflow + +import ( + "context" + "encoding/json" + "fmt" + "os" + "path/filepath" + "slices" + "strings" + + "github.com/entireio/cli/cmd/entire/cli/agent" + "github.com/entireio/cli/cmd/entire/cli/jsonutil" + "github.com/entireio/cli/cmd/entire/cli/paths" +) + +// Ensure IFlowCLIAgent implements HookSupport +var _ agent.HookSupport = (*IFlowCLIAgent)(nil) + +// iFlow hook names - these become subcommands under `entire hooks iflow` +const ( + HookNamePreToolUse = "pre-tool-use" + HookNamePostToolUse = "post-tool-use" + HookNameSetUpEnvironment = "set-up-environment" + HookNameStop = "stop" + HookNameSubagentStop = "subagent-stop" + HookNameSessionStart = "session-start" + HookNameSessionEnd = "session-end" + HookNameUserPromptSubmit = "user-prompt-submit" + HookNameNotification = "notification" +) + +// HookNames returns the hook verbs iFlow CLI supports. +// These become subcommands: entire hooks iflow +func (i *IFlowCLIAgent) HookNames() []string { + return []string{ + HookNamePreToolUse, + HookNamePostToolUse, + HookNameSetUpEnvironment, + HookNameStop, + HookNameSubagentStop, + HookNameSessionStart, + HookNameSessionEnd, + HookNameUserPromptSubmit, + HookNameNotification, + } +} + +// InstallHooks installs iFlow CLI hooks in .iflow/settings.json. +// If force is true, removes existing Entire hooks before installing. +// Returns the number of hooks installed. +func (i *IFlowCLIAgent) InstallHooks(ctx context.Context, localDev bool, force bool) (int, error) { + // Use repo root instead of CWD to find .iflow directory + repoRoot, err := paths.WorktreeRoot(ctx) + if err != nil { + // Fallback to CWD if not in a git repo + repoRoot, err = os.Getwd() + if err != nil { + return 0, fmt.Errorf("failed to get current directory: %w", err) + } + } + + settingsPath := filepath.Join(repoRoot, ".iflow", IFlowSettingsFileName) + + // Read existing settings if they exist + var rawSettings map[string]json.RawMessage + var rawHooks map[string]json.RawMessage + var rawPermissions map[string]json.RawMessage + + existingData, readErr := os.ReadFile(settingsPath) + if readErr == nil { + if err := json.Unmarshal(existingData, &rawSettings); err != nil { + return 0, fmt.Errorf("failed to parse existing settings.json: %w", err) + } + if hooksRaw, ok := rawSettings["hooks"]; ok { + if err := json.Unmarshal(hooksRaw, &rawHooks); err != nil { + return 0, fmt.Errorf("failed to parse hooks in settings.json: %w", err) + } + } + if permRaw, ok := rawSettings["permissions"]; ok { + if err := json.Unmarshal(permRaw, &rawPermissions); err != nil { + return 0, fmt.Errorf("failed to parse permissions in settings.json: %w", err) + } + } + } else { + rawSettings = make(map[string]json.RawMessage) + } + + if rawHooks == nil { + rawHooks = make(map[string]json.RawMessage) + } + if rawPermissions == nil { + rawPermissions = make(map[string]json.RawMessage) + } + + // Parse hook types we need to modify + // All hooks use IFlowHookMatcher format (with hooks array wrapper) per iFlow CLI specification + var preToolUse, postToolUse, sessionStart, userPromptSubmit, notification []IFlowHookMatcher + var setUpEnvironment, stop, subagentStop, sessionEnd []IFlowHookMatcher + + parseHookMatcherType(rawHooks, "PreToolUse", &preToolUse) + parseHookMatcherType(rawHooks, "PostToolUse", &postToolUse) + parseHookMatcherType(rawHooks, "SessionStart", &sessionStart) + parseHookMatcherType(rawHooks, "UserPromptSubmit", &userPromptSubmit) + parseHookMatcherType(rawHooks, "Notification", ¬ification) + parseHookMatcherType(rawHooks, "SetUpEnvironment", &setUpEnvironment) + parseHookMatcherType(rawHooks, "Stop", &stop) + parseHookMatcherType(rawHooks, "SubagentStop", &subagentStop) + parseHookMatcherType(rawHooks, "SessionEnd", &sessionEnd) + + // If force is true, remove all existing Entire hooks first + if force { + preToolUse = removeEntireHooks(preToolUse) + postToolUse = removeEntireHooks(postToolUse) + sessionStart = removeEntireHooks(sessionStart) + userPromptSubmit = removeEntireHooks(userPromptSubmit) + notification = removeEntireHooks(notification) + setUpEnvironment = removeEntireHooks(setUpEnvironment) + stop = removeEntireHooks(stop) + subagentStop = removeEntireHooks(subagentStop) + sessionEnd = removeEntireHooks(sessionEnd) + } + + // Define hook commands + var preToolUseCmd, postToolUseCmd, setUpEnvCmd, stopCmd, subagentStopCmd string + var sessionStartCmd, sessionEndCmd, userPromptSubmitCmd, notificationCmd string + + if localDev { + baseCmd := "go run ${IFLOW_PROJECT_DIR}/cmd/entire/main.go hooks iflow" + preToolUseCmd = baseCmd + " pre-tool-use" + postToolUseCmd = baseCmd + " post-tool-use" + setUpEnvCmd = baseCmd + " set-up-environment" + stopCmd = baseCmd + " stop" + subagentStopCmd = baseCmd + " subagent-stop" + sessionStartCmd = baseCmd + " session-start" + sessionEndCmd = baseCmd + " session-end" + userPromptSubmitCmd = baseCmd + " user-prompt-submit" + notificationCmd = baseCmd + " notification" + } else { + preToolUseCmd = "entire hooks iflow pre-tool-use" + postToolUseCmd = "entire hooks iflow post-tool-use" + setUpEnvCmd = "entire hooks iflow set-up-environment" + stopCmd = "entire hooks iflow stop" + subagentStopCmd = "entire hooks iflow subagent-stop" + sessionStartCmd = "entire hooks iflow session-start" + sessionEndCmd = "entire hooks iflow session-end" + userPromptSubmitCmd = "entire hooks iflow user-prompt-submit" + notificationCmd = "entire hooks iflow notification" + } + + count := 0 + + // Add PreToolUse hook with matcher "*" (all tools) + if !hookMatcherExists(preToolUse, "*", preToolUseCmd) { + preToolUse = addHookMatcher(preToolUse, "*", preToolUseCmd) + count++ + } + + // Add PostToolUse hook with matcher "*" (all tools) + if !hookMatcherExists(postToolUse, "*", postToolUseCmd) { + postToolUse = addHookMatcher(postToolUse, "*", postToolUseCmd) + count++ + } + + // Add SessionStart hook with matcher "startup" + if !hookMatcherExists(sessionStart, "startup", sessionStartCmd) { + sessionStart = addHookMatcher(sessionStart, "startup", sessionStartCmd) + count++ + } + + // Add UserPromptSubmit hook + if !hookMatcherHasCommand(userPromptSubmit, userPromptSubmitCmd) { + userPromptSubmit = append(userPromptSubmit, IFlowHookMatcher{ + Hooks: []IFlowHookEntry{{Type: "command", Command: userPromptSubmitCmd}}, + }) + count++ + } + + // Add Notification hook + if !hookMatcherHasCommand(notification, notificationCmd) { + notification = append(notification, IFlowHookMatcher{ + Hooks: []IFlowHookEntry{{Type: "command", Command: notificationCmd}}, + }) + count++ + } + + // Add SetUpEnvironment hook (no matcher, just hooks array) + if !hookMatcherHasCommand(setUpEnvironment, setUpEnvCmd) { + setUpEnvironment = append(setUpEnvironment, IFlowHookMatcher{ + Hooks: []IFlowHookEntry{{Type: "command", Command: setUpEnvCmd}}, + }) + count++ + } + + // Add Stop hook (no matcher, just hooks array) + if !hookMatcherHasCommand(stop, stopCmd) { + stop = append(stop, IFlowHookMatcher{ + Hooks: []IFlowHookEntry{{Type: "command", Command: stopCmd}}, + }) + count++ + } + + // Add SubagentStop hook (no matcher, just hooks array) + if !hookMatcherHasCommand(subagentStop, subagentStopCmd) { + subagentStop = append(subagentStop, IFlowHookMatcher{ + Hooks: []IFlowHookEntry{{Type: "command", Command: subagentStopCmd}}, + }) + count++ + } + + // Add SessionEnd hook (no matcher, just hooks array) + if !hookMatcherHasCommand(sessionEnd, sessionEndCmd) { + sessionEnd = append(sessionEnd, IFlowHookMatcher{ + Hooks: []IFlowHookEntry{{Type: "command", Command: sessionEndCmd}}, + }) + count++ + } + + // Add permissions.deny rule if not present + permissionsChanged := false + var denyRules []string + if denyRaw, ok := rawPermissions["deny"]; ok { + if err := json.Unmarshal(denyRaw, &denyRules); err != nil { + return 0, fmt.Errorf("failed to parse permissions.deny in settings.json: %w", err) + } + } + if !slices.Contains(denyRules, metadataDenyRule) { + denyRules = append(denyRules, metadataDenyRule) + denyJSON, err := json.Marshal(denyRules) + if err != nil { + return 0, fmt.Errorf("failed to marshal permissions.deny: %w", err) + } + rawPermissions["deny"] = denyJSON + permissionsChanged = true + } + + if count == 0 && !permissionsChanged { + return 0, nil // All hooks and permissions already installed + } + + // Marshal modified hook types back to rawHooks + marshalHookMatcherType(rawHooks, "PreToolUse", preToolUse) + marshalHookMatcherType(rawHooks, "PostToolUse", postToolUse) + marshalHookMatcherType(rawHooks, "SessionStart", sessionStart) + marshalHookMatcherType(rawHooks, "UserPromptSubmit", userPromptSubmit) + marshalHookMatcherType(rawHooks, "Notification", notification) + marshalHookMatcherType(rawHooks, "SetUpEnvironment", setUpEnvironment) + marshalHookMatcherType(rawHooks, "Stop", stop) + marshalHookMatcherType(rawHooks, "SubagentStop", subagentStop) + marshalHookMatcherType(rawHooks, "SessionEnd", sessionEnd) + + // Marshal hooks and update raw settings + hooksJSON, err := json.Marshal(rawHooks) + if err != nil { + return 0, fmt.Errorf("failed to marshal hooks: %w", err) + } + rawSettings["hooks"] = hooksJSON + + // Marshal permissions and update raw settings + permJSON, err := json.Marshal(rawPermissions) + if err != nil { + return 0, fmt.Errorf("failed to marshal permissions: %w", err) + } + rawSettings["permissions"] = permJSON + + // Write back to file + if err := os.MkdirAll(filepath.Dir(settingsPath), 0o750); err != nil { + return 0, fmt.Errorf("failed to create .iflow directory: %w", err) + } + + output, err := jsonutil.MarshalIndentWithNewline(rawSettings, "", " ") + if err != nil { + return 0, fmt.Errorf("failed to marshal settings: %w", err) + } + + if err := os.WriteFile(settingsPath, output, 0o600); err != nil { + return 0, fmt.Errorf("failed to write settings.json: %w", err) + } + + return count, nil +} + +// UninstallHooks removes Entire hooks from iFlow CLI settings. +func (i *IFlowCLIAgent) UninstallHooks(ctx context.Context) error { + // Use repo root to find .iflow directory when run from a subdirectory + repoRoot, err := paths.WorktreeRoot(ctx) + if err != nil { + repoRoot = "." + } + settingsPath := filepath.Join(repoRoot, ".iflow", IFlowSettingsFileName) + data, err := os.ReadFile(settingsPath) + if err != nil { + return nil + } + + var rawSettings map[string]json.RawMessage + if err := json.Unmarshal(data, &rawSettings); err != nil { + return fmt.Errorf("failed to parse settings.json: %w", err) + } + + var rawHooks map[string]json.RawMessage + if hooksRaw, ok := rawSettings["hooks"]; ok { + if err := json.Unmarshal(hooksRaw, &rawHooks); err != nil { + return fmt.Errorf("failed to parse hooks: %w", err) + } + } + if rawHooks == nil { + rawHooks = make(map[string]json.RawMessage) + } + + // Parse and clean all hook types + // All hooks use IFlowHookMatcher format (with hooks array wrapper) per iFlow CLI specification + var preToolUse, postToolUse, sessionStart, userPromptSubmit, notification []IFlowHookMatcher + var setUpEnvironment, stop, subagentStop, sessionEnd []IFlowHookMatcher + + parseHookMatcherType(rawHooks, "PreToolUse", &preToolUse) + parseHookMatcherType(rawHooks, "PostToolUse", &postToolUse) + parseHookMatcherType(rawHooks, "SessionStart", &sessionStart) + parseHookMatcherType(rawHooks, "UserPromptSubmit", &userPromptSubmit) + parseHookMatcherType(rawHooks, "Notification", ¬ification) + parseHookMatcherType(rawHooks, "SetUpEnvironment", &setUpEnvironment) + parseHookMatcherType(rawHooks, "Stop", &stop) + parseHookMatcherType(rawHooks, "SubagentStop", &subagentStop) + parseHookMatcherType(rawHooks, "SessionEnd", &sessionEnd) + + // Remove Entire hooks from all hook types + preToolUse = removeEntireHooks(preToolUse) + postToolUse = removeEntireHooks(postToolUse) + sessionStart = removeEntireHooks(sessionStart) + userPromptSubmit = removeEntireHooks(userPromptSubmit) + notification = removeEntireHooks(notification) + setUpEnvironment = removeEntireHooks(setUpEnvironment) + stop = removeEntireHooks(stop) + subagentStop = removeEntireHooks(subagentStop) + sessionEnd = removeEntireHooks(sessionEnd) + + // Marshal modified hook types back to rawHooks + marshalHookMatcherType(rawHooks, "PreToolUse", preToolUse) + marshalHookMatcherType(rawHooks, "PostToolUse", postToolUse) + marshalHookMatcherType(rawHooks, "SessionStart", sessionStart) + marshalHookMatcherType(rawHooks, "UserPromptSubmit", userPromptSubmit) + marshalHookMatcherType(rawHooks, "Notification", notification) + marshalHookMatcherType(rawHooks, "SetUpEnvironment", setUpEnvironment) + marshalHookMatcherType(rawHooks, "Stop", stop) + marshalHookMatcherType(rawHooks, "SubagentStop", subagentStop) + marshalHookMatcherType(rawHooks, "SessionEnd", sessionEnd) + + // Also remove the metadata deny rule from permissions + var rawPermissions map[string]json.RawMessage + if permRaw, ok := rawSettings["permissions"]; ok { + if err := json.Unmarshal(permRaw, &rawPermissions); err != nil { + rawPermissions = nil + } + } + + if rawPermissions != nil { + if denyRaw, ok := rawPermissions["deny"]; ok { + var denyRules []string + if err := json.Unmarshal(denyRaw, &denyRules); err == nil { + filteredRules := make([]string, 0, len(denyRules)) + for _, rule := range denyRules { + if rule != metadataDenyRule { + filteredRules = append(filteredRules, rule) + } + } + if len(filteredRules) > 0 { + denyJSON, err := json.Marshal(filteredRules) + if err == nil { + rawPermissions["deny"] = denyJSON + } + } else { + delete(rawPermissions, "deny") + } + } + } + + if len(rawPermissions) > 0 { + permJSON, err := json.Marshal(rawPermissions) + if err == nil { + rawSettings["permissions"] = permJSON + } + } else { + delete(rawSettings, "permissions") + } + } + + // Marshal hooks back + if len(rawHooks) > 0 { + hooksJSON, err := json.Marshal(rawHooks) + if err != nil { + return fmt.Errorf("failed to marshal hooks: %w", err) + } + rawSettings["hooks"] = hooksJSON + } else { + delete(rawSettings, "hooks") + } + + // Write back + output, err := jsonutil.MarshalIndentWithNewline(rawSettings, "", " ") + if err != nil { + return fmt.Errorf("failed to marshal settings: %w", err) + } + if err := os.WriteFile(settingsPath, output, 0o600); err != nil { + return fmt.Errorf("failed to write settings.json: %w", err) + } + return nil +} + +// AreHooksInstalled checks if Entire hooks are currently installed. +func (i *IFlowCLIAgent) AreHooksInstalled(ctx context.Context) bool { + repoRoot, err := paths.WorktreeRoot(ctx) + if err != nil { + repoRoot = "." + } + settingsPath := filepath.Join(repoRoot, ".iflow", IFlowSettingsFileName) + data, err := os.ReadFile(settingsPath) + if err != nil { + return false + } + + var settings IFlowSettings + if err := json.Unmarshal(data, &settings); err != nil { + return false + } + + // Check for at least one of our hooks in Stop hook configuration + return hookMatcherHasCommand(settings.Hooks.Stop, "entire hooks iflow stop") || + hookMatcherHasCommand(settings.Hooks.Stop, "go run ${IFLOW_PROJECT_DIR}/cmd/entire/main.go hooks iflow stop") +} + +// Helper functions for hook management + +func parseHookMatcherType(rawHooks map[string]json.RawMessage, hookType string, target *[]IFlowHookMatcher) { + if data, ok := rawHooks[hookType]; ok { + json.Unmarshal(data, target) + } +} + +func marshalHookMatcherType(rawHooks map[string]json.RawMessage, hookType string, matchers []IFlowHookMatcher) { + if len(matchers) == 0 { + delete(rawHooks, hookType) + return + } + data, err := json.Marshal(matchers) + if err != nil { + return + } + rawHooks[hookType] = data +} + +func isEntireHook(command string) bool { + for _, prefix := range entireHookPrefixes { + if strings.HasPrefix(command, prefix) { + return true + } + } + return false +} + +func removeEntireHooks(matchers []IFlowHookMatcher) []IFlowHookMatcher { + result := make([]IFlowHookMatcher, 0, len(matchers)) + for _, matcher := range matchers { + filteredHooks := make([]IFlowHookEntry, 0, len(matcher.Hooks)) + for _, hook := range matcher.Hooks { + if !isEntireHook(hook.Command) { + filteredHooks = append(filteredHooks, hook) + } + } + if len(filteredHooks) > 0 { + matcher.Hooks = filteredHooks + result = append(result, matcher) + } + } + return result +} + +func hookMatcherExists(matchers []IFlowHookMatcher, matcherPattern, command string) bool { + for _, matcher := range matchers { + if matcher.Matcher == matcherPattern { + for _, hook := range matcher.Hooks { + if hook.Command == command { + return true + } + } + } + } + return false +} + +func hookMatcherHasCommand(matchers []IFlowHookMatcher, command string) bool { + for _, matcher := range matchers { + for _, hook := range matcher.Hooks { + if hook.Command == command { + return true + } + } + } + return false +} + +func addHookMatcher(matchers []IFlowHookMatcher, matcherPattern, command string) []IFlowHookMatcher { + entry := IFlowHookEntry{ + Type: "command", + Command: command, + } + + for i, matcher := range matchers { + if matcher.Matcher == matcherPattern { + matchers[i].Hooks = append(matchers[i].Hooks, entry) + return matchers + } + } + + return append(matchers, IFlowHookMatcher{ + Matcher: matcherPattern, + Hooks: []IFlowHookEntry{entry}, + }) +} diff --git a/cmd/entire/cli/agent/iflow/hooks_test.go b/cmd/entire/cli/agent/iflow/hooks_test.go new file mode 100644 index 000000000..e8cbd6b71 --- /dev/null +++ b/cmd/entire/cli/agent/iflow/hooks_test.go @@ -0,0 +1,335 @@ +package iflow + +import ( + "context" + "encoding/json" + "os" + "path/filepath" + "testing" +) + +func TestInstallHooks(t *testing.T) { + dir := t.TempDir() + t.Chdir(dir) + + // Create .iflow directory + if err := os.MkdirAll(filepath.Join(dir, ".iflow"), 0o755); err != nil { + t.Fatal(err) + } + + ag := NewIFlowCLIAgent().(*IFlowCLIAgent) + ctx := context.Background() + + // Install hooks + count, err := ag.InstallHooks(ctx, false, false) + if err != nil { + t.Fatalf("InstallHooks failed: %v", err) + } + + if count == 0 { + t.Error("Expected hooks to be installed, got 0") + } + + // Verify settings.json was created + settingsPath := filepath.Join(dir, ".iflow", IFlowSettingsFileName) + data, err := os.ReadFile(settingsPath) + if err != nil { + t.Fatalf("Failed to read settings.json: %v", err) + } + + var settings IFlowSettings + if err := json.Unmarshal(data, &settings); err != nil { + t.Fatalf("Failed to unmarshal settings.json: %v", err) + } + + // Check that hooks were installed + if len(settings.Hooks.Stop) == 0 { + t.Error("Expected Stop hooks to be installed") + } + if len(settings.Hooks.SessionStart) == 0 { + t.Error("Expected SessionStart hooks to be installed") + } + if len(settings.Hooks.PreToolUse) == 0 { + t.Error("Expected PreToolUse hooks to be installed") + } + + // Check permissions + if len(settings.Permissions.Deny) == 0 { + t.Error("Expected permissions.deny to be set") + } +} + +func TestInstallHooks_Idempotent(t *testing.T) { + dir := t.TempDir() + t.Chdir(dir) + + if err := os.MkdirAll(filepath.Join(dir, ".iflow"), 0o755); err != nil { + t.Fatal(err) + } + + ag := NewIFlowCLIAgent().(*IFlowCLIAgent) + ctx := context.Background() + + // Install hooks first time + count1, err := ag.InstallHooks(ctx, false, false) + if err != nil { + t.Fatalf("First InstallHooks failed: %v", err) + } + if count1 == 0 { + t.Error("Expected hooks to be installed on first run") + } + + // Install hooks second time (should be idempotent) + count2, err := ag.InstallHooks(ctx, false, false) + if err != nil { + t.Fatalf("Second InstallHooks failed: %v", err) + } + if count2 != 0 { + t.Errorf("Expected 0 hooks on second install, got %d", count2) + } +} + +func TestInstallHooks_Force(t *testing.T) { + dir := t.TempDir() + t.Chdir(dir) + + if err := os.MkdirAll(filepath.Join(dir, ".iflow"), 0o755); err != nil { + t.Fatal(err) + } + + ag := NewIFlowCLIAgent().(*IFlowCLIAgent) + ctx := context.Background() + + // Install hooks first + _, err := ag.InstallHooks(ctx, false, false) + if err != nil { + t.Fatalf("First InstallHooks failed: %v", err) + } + + // Force reinstall + count, err := ag.InstallHooks(ctx, false, true) + if err != nil { + t.Fatalf("Force InstallHooks failed: %v", err) + } + if count == 0 { + t.Error("Expected hooks to be reinstalled with force=true") + } +} + +func TestUninstallHooks(t *testing.T) { + dir := t.TempDir() + t.Chdir(dir) + + if err := os.MkdirAll(filepath.Join(dir, ".iflow"), 0o755); err != nil { + t.Fatal(err) + } + + ag := NewIFlowCLIAgent().(*IFlowCLIAgent) + ctx := context.Background() + + // Install hooks + _, err := ag.InstallHooks(ctx, false, false) + if err != nil { + t.Fatalf("InstallHooks failed: %v", err) + } + + // Verify hooks are installed + if !ag.AreHooksInstalled(ctx) { + t.Error("Expected hooks to be installed") + } + + // Uninstall hooks + if err := ag.UninstallHooks(ctx); err != nil { + t.Fatalf("UninstallHooks failed: %v", err) + } + + // Verify hooks are not installed + if ag.AreHooksInstalled(ctx) { + t.Error("Expected hooks to be uninstalled") + } +} + +func TestAreHooksInstalled(t *testing.T) { + dir := t.TempDir() + t.Chdir(dir) + + ag := NewIFlowCLIAgent().(*IFlowCLIAgent) + ctx := context.Background() + + // Initially no hooks + if ag.AreHooksInstalled(ctx) { + t.Error("Expected no hooks initially") + } + + // Create .iflow directory and install hooks + if err := os.MkdirAll(filepath.Join(dir, ".iflow"), 0o755); err != nil { + t.Fatal(err) + } + + _, err := ag.InstallHooks(ctx, false, false) + if err != nil { + t.Fatalf("InstallHooks failed: %v", err) + } + + // Now hooks should be installed + if !ag.AreHooksInstalled(ctx) { + t.Error("Expected hooks to be detected") + } +} + +func TestHookNames(t *testing.T) { + ag := NewIFlowCLIAgent().(*IFlowCLIAgent) + names := ag.HookNames() + + expected := []string{ + HookNamePreToolUse, + HookNamePostToolUse, + HookNameSetUpEnvironment, + HookNameStop, + HookNameSubagentStop, + HookNameSessionStart, + HookNameSessionEnd, + HookNameUserPromptSubmit, + HookNameNotification, + } + + if len(names) != len(expected) { + t.Errorf("Expected %d hook names, got %d", len(expected), len(names)) + } + + for i, name := range expected { + if i >= len(names) || names[i] != name { + t.Errorf("Expected hook name %q at position %d, got %q", name, i, names[i]) + } + } +} + +func TestIsEntireHook(t *testing.T) { + tests := []struct { + command string + expected bool + }{ + {"entire hooks iflow stop", true}, + {"go run ${IFLOW_PROJECT_DIR}/cmd/entire/main.go hooks iflow stop", true}, + {"some other command", false}, + {"", false}, + } + + for _, tt := range tests { + t.Run(tt.command, func(t *testing.T) { + result := isEntireHook(tt.command) + if result != tt.expected { + t.Errorf("isEntireHook(%q) = %v, want %v", tt.command, result, tt.expected) + } + }) + } +} + +func TestAddHookMatcher(t *testing.T) { + tests := []struct { + name string + initial []IFlowHookMatcher + matcherPattern string + command string + expectedLen int + }{ + { + name: "add to empty", + initial: []IFlowHookMatcher{}, + matcherPattern: "*", + command: "test command", + expectedLen: 1, + }, + { + name: "add to existing matcher", + initial: []IFlowHookMatcher{ + {Matcher: "*", Hooks: []IFlowHookEntry{{Type: "command", Command: "existing"}}}, + }, + matcherPattern: "*", + command: "new command", + expectedLen: 1, + }, + { + name: "add new matcher", + initial: []IFlowHookMatcher{ + {Matcher: "Edit", Hooks: []IFlowHookEntry{{Type: "command", Command: "edit hook"}}}, + }, + matcherPattern: "Write", + command: "write hook", + expectedLen: 2, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := addHookMatcher(tt.initial, tt.matcherPattern, tt.command) + if len(result) != tt.expectedLen { + t.Errorf("Expected %d matchers, got %d", tt.expectedLen, len(result)) + } + }) + } +} + +func TestHookMatcherExists(t *testing.T) { + matchers := []IFlowHookMatcher{ + { + Matcher: "*", + Hooks: []IFlowHookEntry{ + {Type: "command", Command: "entire hooks iflow stop"}, + }, + }, + } + + tests := []struct { + matcherPattern string + command string + expected bool + }{ + {"*", "entire hooks iflow stop", true}, + {"*", "other command", false}, + {"Edit", "entire hooks iflow stop", false}, + } + + for _, tt := range tests { + t.Run(tt.command, func(t *testing.T) { + result := hookMatcherExists(matchers, tt.matcherPattern, tt.command) + if result != tt.expected { + t.Errorf("hookMatcherExists(%q, %q) = %v, want %v", tt.matcherPattern, tt.command, result, tt.expected) + } + }) + } +} + +func TestRemoveEntireHooks(t *testing.T) { + matchers := []IFlowHookMatcher{ + { + Matcher: "*", + Hooks: []IFlowHookEntry{ + {Type: "command", Command: "entire hooks iflow stop"}, + {Type: "command", Command: "other hook"}, + }, + }, + { + Matcher: "Edit", + Hooks: []IFlowHookEntry{ + {Type: "command", Command: "entire hooks iflow pre-tool-use"}, + }, + }, + } + + result := removeEntireHooks(matchers) + + // Should have 1 matcher (Edit matcher removed because all hooks were Entire hooks) + if len(result) != 1 { + t.Errorf("Expected 1 matcher after removal, got %d", len(result)) + } + + // The remaining matcher should have 1 hook (other hook) + if len(result[0].Hooks) != 1 { + t.Errorf("Expected 1 hook in remaining matcher, got %d", len(result[0].Hooks)) + } + + if result[0].Hooks[0].Command != "other hook" { + t.Errorf("Expected 'other hook', got %q", result[0].Hooks[0].Command) + } +} diff --git a/cmd/entire/cli/agent/iflow/iflow.go b/cmd/entire/cli/agent/iflow/iflow.go new file mode 100644 index 000000000..213a79e0b --- /dev/null +++ b/cmd/entire/cli/agent/iflow/iflow.go @@ -0,0 +1,307 @@ +// Package iflow implements the Agent interface for iFlow CLI. +// iFlow CLI is Alibaba's AI coding assistant with a hooks-based event system. +package iflow + +import ( + "bufio" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "os" + "path/filepath" + "regexp" + "time" + + "github.com/entireio/cli/cmd/entire/cli/agent" + "github.com/entireio/cli/cmd/entire/cli/agent/types" + "github.com/entireio/cli/cmd/entire/cli/paths" +) + +//nolint:gochecknoinits // Agent self-registration is the intended pattern +func init() { + agent.Register(agent.AgentNameIFlow, NewIFlowCLIAgent) +} + +// IFlowCLIAgent implements the Agent interface for iFlow CLI. +type IFlowCLIAgent struct{} + +// NewIFlowCLIAgent creates a new iFlow CLI agent instance. +func NewIFlowCLIAgent() agent.Agent { + return &IFlowCLIAgent{} +} + +// Name returns the agent registry key. +func (i *IFlowCLIAgent) Name() types.AgentName { + return agent.AgentNameIFlow +} + +// Type returns the agent type identifier. +func (i *IFlowCLIAgent) Type() types.AgentType { + return agent.AgentTypeIFlow +} + +// Description returns a human-readable description. +func (i *IFlowCLIAgent) Description() string { + return "iFlow CLI - Alibaba's AI coding assistant" +} + +// IsPreview returns whether the agent integration is in preview. +func (i *IFlowCLIAgent) IsPreview() bool { + return true +} + +// DetectPresence checks if iFlow CLI is configured in the repository. +func (i *IFlowCLIAgent) DetectPresence(ctx context.Context) (bool, error) { + // Get worktree root to check for .iflow directory + repoRoot, err := paths.WorktreeRoot(ctx) + if err != nil { + // Not in a git repo, fall back to CWD-relative check + repoRoot = "." + } + + // Check for .iflow directory + iflowDir := filepath.Join(repoRoot, ".iflow") + if _, err := os.Stat(iflowDir); err == nil { + return true, nil + } + // Check for .iflow/settings.json + settingsFile := filepath.Join(iflowDir, "settings.json") + if _, err := os.Stat(settingsFile); err == nil { + return true, nil + } + return false, nil +} + +// ProtectedDirs returns directories that iFlow uses for config/state. +func (i *IFlowCLIAgent) ProtectedDirs() []string { + return []string{".iflow"} +} + +// GetSessionID extracts the session ID from hook input. +func (i *IFlowCLIAgent) GetSessionID(input *agent.HookInput) string { + return input.SessionID +} + +// GetSessionDir returns the directory where iFlow stores session transcripts. +func (i *IFlowCLIAgent) GetSessionDir(repoPath string) (string, error) { + // Check for test environment override + if override := os.Getenv("ENTIRE_TEST_IFLOW_PROJECT_DIR"); override != "" { + return override, nil + } + + homeDir, err := os.UserHomeDir() + if err != nil { + return "", fmt.Errorf("failed to get home directory: %w", err) + } + + projectDir := SanitizePathForIFlow(repoPath) + return filepath.Join(homeDir, ".iflow", "projects", projectDir), nil +} + +// ResolveSessionFile returns the path to an iFlow session file. +// iFlow names files as .jsonl +func (i *IFlowCLIAgent) ResolveSessionFile(sessionDir, agentSessionID string) string { + return filepath.Join(sessionDir, agentSessionID+".jsonl") +} + +// ReadSession reads a session from iFlow's storage (JSONL transcript file). +func (i *IFlowCLIAgent) ReadSession(input *agent.HookInput) (*agent.AgentSession, error) { + if input.SessionRef == "" { + return nil, errors.New("session reference (transcript path) is required") + } + + // Read the raw JSONL file + data, err := os.ReadFile(input.SessionRef) + if err != nil { + return nil, fmt.Errorf("failed to read transcript: %w", err) + } + + // Parse to extract computed fields + lines, err := ParseTranscriptFromBytes(data) + if err != nil { + return nil, fmt.Errorf("failed to parse transcript: %w", err) + } + + return &agent.AgentSession{ + SessionID: input.SessionID, + AgentName: i.Name(), + SessionRef: input.SessionRef, + StartTime: time.Now(), + NativeData: data, + ModifiedFiles: ExtractModifiedFiles(lines), + }, nil +} + +// WriteSession writes a session to iFlow's storage (JSONL transcript file). +func (i *IFlowCLIAgent) WriteSession(_ context.Context, session *agent.AgentSession) error { + if session == nil { + return errors.New("session is nil") + } + + // Verify this session belongs to iFlow + if session.AgentName != "" && session.AgentName != i.Name() { + return fmt.Errorf("session belongs to agent %q, not %q", session.AgentName, i.Name()) + } + + if session.SessionRef == "" { + return errors.New("session reference (transcript path) is required") + } + + if len(session.NativeData) == 0 { + return errors.New("session has no native data to write") + } + + // Write the raw JSONL data + if err := os.WriteFile(session.SessionRef, session.NativeData, 0o600); err != nil { + return fmt.Errorf("failed to write transcript: %w", err) + } + + return nil +} + +// FormatResumeCommand returns the command to resume an iFlow CLI session. +func (i *IFlowCLIAgent) FormatResumeCommand(sessionID string) string { + return "iflow -r " + sessionID +} + +// ReadTranscript reads the raw JSONL transcript bytes for a session. +func (i *IFlowCLIAgent) ReadTranscript(sessionRef string) ([]byte, error) { + data, err := os.ReadFile(sessionRef) + if err != nil { + return nil, fmt.Errorf("failed to read transcript: %w", err) + } + return data, nil +} + +// ChunkTranscript splits a JSONL transcript at line boundaries. +func (i *IFlowCLIAgent) ChunkTranscript(_ context.Context, content []byte, maxSize int) ([][]byte, error) { + chunks, err := agent.ChunkJSONL(content, maxSize) + if err != nil { + return nil, fmt.Errorf("failed to chunk JSONL transcript: %w", err) + } + return chunks, nil +} + +// ReassembleTranscript concatenates JSONL chunks with newlines. +func (i *IFlowCLIAgent) ReassembleTranscript(chunks [][]byte) ([]byte, error) { + return agent.ReassembleJSONL(chunks), nil +} + +// GetTranscriptPosition returns the current line count of an iFlow transcript. +func (i *IFlowCLIAgent) GetTranscriptPosition(path string) (int, error) { + if path == "" { + return 0, nil + } + + file, err := os.Open(path) + if err != nil { + if os.IsNotExist(err) { + return 0, nil + } + return 0, fmt.Errorf("failed to open transcript file: %w", err) + } + defer file.Close() + + reader := bufio.NewReader(file) + lineCount := 0 + + for { + line, err := reader.ReadBytes('\n') + if err != nil { + if err == io.EOF { + if len(line) > 0 { + lineCount++ + } + break + } + return 0, fmt.Errorf("failed to read transcript: %w", err) + } + lineCount++ + } + + return lineCount, nil +} + +// ExtractModifiedFilesFromOffset extracts files modified since a given line number. +func (i *IFlowCLIAgent) ExtractModifiedFilesFromOffset(path string, startOffset int) (files []string, currentPosition int, err error) { + if path == "" { + return nil, 0, nil + } + + file, openErr := os.Open(path) + if openErr != nil { + return nil, 0, fmt.Errorf("failed to open transcript file: %w", openErr) + } + defer file.Close() + + reader := bufio.NewReader(file) + var lines []TranscriptLine + lineNum := 0 + + for { + lineData, readErr := reader.ReadBytes('\n') + if readErr != nil && readErr != io.EOF { + return nil, 0, fmt.Errorf("failed to read transcript: %w", readErr) + } + + if len(lineData) > 0 { + lineNum++ + if lineNum > startOffset { + var line TranscriptLine + if parseErr := json.Unmarshal(lineData, &line); parseErr == nil { + lines = append(lines, line) + } + } + } + + if readErr == io.EOF { + break + } + } + + return ExtractModifiedFiles(lines), lineNum, nil +} + +// SanitizePathForIFlow converts a path to iFlow's project directory format. +// iFlow replaces any non-alphanumeric character with a dash. +var nonAlphanumericRegex = regexp.MustCompile(`[^a-zA-Z0-9]`) + +func SanitizePathForIFlow(path string) string { + return nonAlphanumericRegex.ReplaceAllString(path, "-") +} + +// CalculateTokenUsage computes token usage from the transcript starting at the given line offset. +// iFlow transcripts are JSONL format where each line may contain token usage data. +// If token data is not available in the transcript, returns empty TokenUsage. +func (i *IFlowCLIAgent) CalculateTokenUsage(transcriptData []byte, fromOffset int) (*agent.TokenUsage, error) { + lines, err := ParseTranscriptFromBytes(transcriptData) + if err != nil { + return &agent.TokenUsage{}, fmt.Errorf("failed to parse transcript for token usage: %w", err) + } + + usage := &agent.TokenUsage{} + + // Skip lines before fromOffset + for idx := fromOffset; idx < len(lines); idx++ { + line := lines[idx] + + // Only count tokens from assistant messages (type "assistant" or "gemini") + // iFlow may use different message types + if line.Type != "assistant" && line.Type != "gemini" && line.Type != "ai" { + continue + } + + if line.Tokens == nil { + continue + } + + usage.APICallCount++ + usage.InputTokens += line.Tokens.Input + usage.OutputTokens += line.Tokens.Output + usage.CacheReadTokens += line.Tokens.Cached + } + + return usage, nil +} diff --git a/cmd/entire/cli/agent/iflow/iflow_test.go b/cmd/entire/cli/agent/iflow/iflow_test.go new file mode 100644 index 000000000..5c92ec40f --- /dev/null +++ b/cmd/entire/cli/agent/iflow/iflow_test.go @@ -0,0 +1,312 @@ +package iflow + +import ( + "context" + "os" + "path/filepath" + "testing" + + "github.com/entireio/cli/cmd/entire/cli/agent" +) + +func TestIFlowCLIAgent_Name(t *testing.T) { + ag := NewIFlowCLIAgent() + if ag.Name() != agent.AgentNameIFlow { + t.Errorf("Expected name %q, got %q", agent.AgentNameIFlow, ag.Name()) + } +} + +func TestIFlowCLIAgent_Type(t *testing.T) { + ag := NewIFlowCLIAgent() + if ag.Type() != agent.AgentTypeIFlow { + t.Errorf("Expected type %q, got %q", agent.AgentTypeIFlow, ag.Type()) + } +} + +func TestIFlowCLIAgent_Description(t *testing.T) { + ag := NewIFlowCLIAgent() + expected := "iFlow CLI - Alibaba's AI coding assistant" + if ag.Description() != expected { + t.Errorf("Expected description %q, got %q", expected, ag.Description()) + } +} + +func TestIFlowCLIAgent_IsPreview(t *testing.T) { + ag := NewIFlowCLIAgent() + if !ag.IsPreview() { + t.Error("Expected IsPreview to return true") + } +} + +func TestIFlowCLIAgent_ProtectedDirs(t *testing.T) { + ag := NewIFlowCLIAgent() + dirs := ag.ProtectedDirs() + if len(dirs) != 1 || dirs[0] != ".iflow" { + t.Errorf("Expected protected dirs [.iflow], got %v", dirs) + } +} + +func TestIFlowCLIAgent_DetectPresence(t *testing.T) { + tests := []struct { + name string + setup func(t *testing.T) string + expected bool + }{ + { + name: "detects .iflow directory", + setup: func(t *testing.T) string { + dir := t.TempDir() + if err := os.MkdirAll(filepath.Join(dir, ".iflow"), 0o755); err != nil { + t.Fatal(err) + } + return dir + }, + expected: true, + }, + { + name: "detects settings.json", + setup: func(t *testing.T) string { + dir := t.TempDir() + if err := os.MkdirAll(filepath.Join(dir, ".iflow"), 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(dir, ".iflow", "settings.json"), []byte("{}"), 0o600); err != nil { + t.Fatal(err) + } + return dir + }, + expected: true, + }, + { + name: "no iFlow config", + setup: func(t *testing.T) string { + return t.TempDir() + }, + expected: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + dir := tt.setup(t) + t.Chdir(dir) + + ag := NewIFlowCLIAgent() + present, err := ag.DetectPresence(context.Background()) + if err != nil { + t.Fatalf("DetectPresence failed: %v", err) + } + if present != tt.expected { + t.Errorf("Expected DetectPresence=%v, got %v", tt.expected, present) + } + }) + } +} + +func TestIFlowCLIAgent_FormatResumeCommand(t *testing.T) { + ag := NewIFlowCLIAgent() + sessionID := "test-session-123" + cmd := ag.FormatResumeCommand(sessionID) + expected := "iflow -r test-session-123" + if cmd != expected { + t.Errorf("Expected resume command %q, got %q", expected, cmd) + } +} + +func TestSanitizePathForIFlow(t *testing.T) { + tests := []struct { + input string + expected string + }{ + {"/home/user/project", "-home-user-project"}, + {"my-project", "my-project"}, + {"my_project", "my-project"}, + {"my.project", "my-project"}, + {"my@project", "my-project"}, + } + + for _, tt := range tests { + t.Run(tt.input, func(t *testing.T) { + result := SanitizePathForIFlow(tt.input) + if result != tt.expected { + t.Errorf("SanitizePathForIFlow(%q) = %q, want %q", tt.input, result, tt.expected) + } + }) + } +} + +func TestIFlowCLIAgent_ResolveSessionFile(t *testing.T) { + ag := NewIFlowCLIAgent() + sessionDir := "/home/user/.iflow/projects/my-project" + sessionID := "test-session" + + path := ag.ResolveSessionFile(sessionDir, sessionID) + expected := filepath.Join(sessionDir, "test-session.jsonl") + if path != expected { + t.Errorf("Expected session file path %q, got %q", expected, path) + } +} + +func TestIFlowCLIAgent_GetSessionID(t *testing.T) { + ag := NewIFlowCLIAgent() + input := &agent.HookInput{ + SessionID: "test-session-id", + } + + sessionID := ag.GetSessionID(input) + if sessionID != "test-session-id" { + t.Errorf("Expected session ID %q, got %q", "test-session-id", sessionID) + } +} + +// Test interface compliance +func TestIFlowCLIAgent_ImplementsAgent(t *testing.T) { + var _ agent.Agent = (*IFlowCLIAgent)(nil) +} + +func TestIFlowCLIAgent_ImplementsHookSupport(t *testing.T) { + var _ agent.HookSupport = (*IFlowCLIAgent)(nil) +} + +func TestIFlowCLIAgent_ImplementsTranscriptAnalyzer(t *testing.T) { + var _ agent.TranscriptAnalyzer = (*IFlowCLIAgent)(nil) +} + +func TestIFlowCLIAgent_ImplementsHookResponseWriter(t *testing.T) { + var _ agent.HookResponseWriter = (*IFlowCLIAgent)(nil) +} + +// Ensure agent is registered +func TestIFlowCLIAgent_Registered(t *testing.T) { + ag, err := agent.Get(agent.AgentNameIFlow) + if err != nil { + t.Fatalf("iFlow agent not registered: %v", err) + } + if ag == nil { + t.Fatal("iFlow agent is nil") + } + if ag.Name() != agent.AgentNameIFlow { + t.Errorf("Expected name %q, got %q", agent.AgentNameIFlow, ag.Name()) + } +} + +// Test GetByAgentType +func TestGetByAgentType_IFlow(t *testing.T) { + ag, err := agent.GetByAgentType(agent.AgentTypeIFlow) + if err != nil { + t.Fatalf("Failed to get iFlow agent by type: %v", err) + } + if ag.Type() != agent.AgentTypeIFlow { + t.Errorf("Expected type %q, got %q", agent.AgentTypeIFlow, ag.Type()) + } +} + +func TestIFlowCLIAgent_ImplementsTokenCalculator(t *testing.T) { + var _ agent.TokenCalculator = (*IFlowCLIAgent)(nil) +} + +func TestCalculateTokenUsage(t *testing.T) { + ag := NewIFlowCLIAgent().(*IFlowCLIAgent) + + tests := []struct { + name string + transcript string + fromOffset int + expectUsage *agent.TokenUsage + expectErr bool + }{ + { + name: "empty transcript", + transcript: "", + fromOffset: 0, + expectUsage: &agent.TokenUsage{}, + expectErr: false, + }, + { + name: "transcript with tokens", + transcript: `{"type":"user","timestamp":"2024-01-01T00:00:00Z","message":"hello"} +{"type":"assistant","timestamp":"2024-01-01T00:00:01Z","tokens":{"input":100,"output":50,"cached":10}} +{"type":"assistant","timestamp":"2024-01-01T00:00:02Z","tokens":{"input":200,"output":75,"cached":20}}`, + fromOffset: 0, + expectUsage: &agent.TokenUsage{ + InputTokens: 300, + OutputTokens: 125, + CacheReadTokens: 30, + APICallCount: 2, + }, + expectErr: false, + }, + { + name: "transcript with fromOffset", + transcript: `{"type":"assistant","timestamp":"2024-01-01T00:00:01Z","tokens":{"input":100,"output":50}} +{"type":"assistant","timestamp":"2024-01-01T00:00:02Z","tokens":{"input":200,"output":75}}`, + fromOffset: 1, + expectUsage: &agent.TokenUsage{ + InputTokens: 200, + OutputTokens: 75, + APICallCount: 1, + }, + expectErr: false, + }, + { + name: "transcript without tokens", + transcript: `{"type":"user","timestamp":"2024-01-01T00:00:00Z","message":"hello"} +{"type":"assistant","timestamp":"2024-01-01T00:00:01Z","message":"hi there"}`, + fromOffset: 0, + expectUsage: &agent.TokenUsage{}, + expectErr: false, + }, + { + name: "mixed message types", + transcript: `{"type":"user","timestamp":"2024-01-01T00:00:00Z","message":"hello","tokens":{"input":50}} +{"type":"assistant","timestamp":"2024-01-01T00:00:01Z","tokens":{"input":100,"output":50}} +{"type":"ai","timestamp":"2024-01-01T00:00:02Z","tokens":{"input":200,"output":75}}`, + fromOffset: 0, + expectUsage: &agent.TokenUsage{ + InputTokens: 300, + OutputTokens: 125, + APICallCount: 2, + }, + expectErr: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + usage, err := ag.CalculateTokenUsage([]byte(tt.transcript), tt.fromOffset) + if (err != nil) != tt.expectErr { + t.Errorf("CalculateTokenUsage() error = %v, expectErr %v", err, tt.expectErr) + return + } + if usage.InputTokens != tt.expectUsage.InputTokens { + t.Errorf("InputTokens = %d, want %d", usage.InputTokens, tt.expectUsage.InputTokens) + } + if usage.OutputTokens != tt.expectUsage.OutputTokens { + t.Errorf("OutputTokens = %d, want %d", usage.OutputTokens, tt.expectUsage.OutputTokens) + } + if usage.CacheReadTokens != tt.expectUsage.CacheReadTokens { + t.Errorf("CacheReadTokens = %d, want %d", usage.CacheReadTokens, tt.expectUsage.CacheReadTokens) + } + if usage.APICallCount != tt.expectUsage.APICallCount { + t.Errorf("APICallCount = %d, want %d", usage.APICallCount, tt.expectUsage.APICallCount) + } + }) + } +} + +func TestCalculateTokenUsage_InvalidJSON(t *testing.T) { + ag := NewIFlowCLIAgent().(*IFlowCLIAgent) + + // Invalid JSON should not cause panic, but return empty usage + transcript := `{invalid json` + usage, err := ag.CalculateTokenUsage([]byte(transcript), 0) + + // The current implementation should handle this gracefully + // ParseTranscriptFromBytes skips malformed lines + if err != nil { + t.Errorf("Expected no error for malformed JSON, got %v", err) + } + if usage == nil { + t.Error("Expected non-nil usage") + } +} diff --git a/cmd/entire/cli/agent/iflow/lifecycle.go b/cmd/entire/cli/agent/iflow/lifecycle.go new file mode 100644 index 000000000..1a3c6526b --- /dev/null +++ b/cmd/entire/cli/agent/iflow/lifecycle.go @@ -0,0 +1,282 @@ +package iflow + +import ( + "context" + "encoding/json" + "fmt" + "io" + "os" + "time" + + "github.com/entireio/cli/cmd/entire/cli/agent" + "github.com/entireio/cli/cmd/entire/cli/paths" +) + +// Ensure IFlowCLIAgent implements required interfaces +var ( + _ agent.TranscriptAnalyzer = (*IFlowCLIAgent)(nil) + _ agent.HookResponseWriter = (*IFlowCLIAgent)(nil) + _ agent.TokenCalculator = (*IFlowCLIAgent)(nil) +) + +// WriteHookResponse outputs a JSON hook response to stdout. +// iFlow CLI can read this JSON and display messages to the user. +func (i *IFlowCLIAgent) WriteHookResponse(message string) error { + resp := struct { + SystemMessage string `json:"systemMessage,omitempty"` + }{SystemMessage: message} + if err := json.NewEncoder(os.Stdout).Encode(resp); err != nil { + return fmt.Errorf("failed to encode hook response: %w", err) + } + return nil +} + +// ParseHookEvent translates an iFlow CLI hook into a normalized lifecycle Event. +// Returns nil if the hook has no lifecycle significance. +func (i *IFlowCLIAgent) ParseHookEvent(ctx context.Context, hookName string, stdin io.Reader) (*agent.Event, error) { + switch hookName { + case HookNameSessionStart: + return i.parseSessionStart(ctx, stdin) + case HookNameUserPromptSubmit: + return i.parseTurnStart(ctx, stdin) + case HookNamePreToolUse: + return i.parsePreToolUse(stdin) + case HookNamePostToolUse: + return i.parsePostToolUse(stdin) + case HookNameStop: + return i.parseStop(ctx, stdin) + case HookNameSessionEnd: + return i.parseSessionEnd(ctx, stdin) + case HookNameSubagentStop: + return i.parseSubagentStop(ctx, stdin) + case HookNameSetUpEnvironment, HookNameNotification: + // These hooks don't have lifecycle significance for Entire + return nil, nil + default: + return nil, nil + } +} + +// --- Internal hook parsing functions --- + +// computeTranscriptPath returns the transcript path if provided, or computes it from session_id. +// iFlow CLI may not always provide transcript_path in hook input, so we compute it as a fallback. +func (i *IFlowCLIAgent) computeTranscriptPath(ctx context.Context, sessionID, providedPath string) string { + // If transcript_path is provided, use it directly + if providedPath != "" { + return providedPath + } + + // Fallback: compute from session_id + // Path pattern: ~/.iflow/projects//.jsonl + if sessionID == "" { + return "" + } + + repoPath, err := paths.WorktreeRoot(ctx) + if err != nil { + return "" + } + + sessionDir, err := i.GetSessionDir(repoPath) + if err != nil { + return "" + } + + return i.ResolveSessionFile(sessionDir, sessionID) +} + +func (i *IFlowCLIAgent) parseSessionStart(ctx context.Context, stdin io.Reader) (*agent.Event, error) { + var input SessionStartInput + if err := json.NewDecoder(stdin).Decode(&input); err != nil { + return nil, fmt.Errorf("failed to decode session start input: %w", err) + } + + // Override with environment variables if present (iFlow CLI may use env vars) + sessionID := input.SessionID + if envSessionID := os.Getenv(EnvIFlowSessionID); envSessionID != "" { + sessionID = envSessionID + } + transcriptPath := input.TranscriptPath + if envTranscriptPath := os.Getenv(EnvIFlowTranscriptPath); envTranscriptPath != "" { + transcriptPath = envTranscriptPath + } + + // Fallback: compute transcript path from session_id if not provided + transcriptPath = i.computeTranscriptPath(ctx, sessionID, transcriptPath) + + event := &agent.Event{ + Type: agent.SessionStart, + SessionID: sessionID, + SessionRef: transcriptPath, + Timestamp: time.Now(), + Metadata: make(map[string]string), + } + + if input.Model != "" { + event.Model = input.Model + } + + if input.Source != "" { + event.Metadata["session_source"] = input.Source + } + + return event, nil +} + +func (i *IFlowCLIAgent) parseTurnStart(ctx context.Context, stdin io.Reader) (*agent.Event, error) { + var input UserPromptSubmitInput + if err := json.NewDecoder(stdin).Decode(&input); err != nil { + return nil, fmt.Errorf("failed to decode user prompt submit input: %w", err) + } + + // Override with environment variables if present (iFlow CLI may use env vars) + sessionID := input.SessionID + if envSessionID := os.Getenv(EnvIFlowSessionID); envSessionID != "" { + sessionID = envSessionID + } + transcriptPath := input.TranscriptPath + if envTranscriptPath := os.Getenv(EnvIFlowTranscriptPath); envTranscriptPath != "" { + transcriptPath = envTranscriptPath + } + prompt := input.Prompt + if envPrompt := os.Getenv(EnvIFlowUserPrompt); envPrompt != "" { + prompt = envPrompt + } + + // Fallback: compute transcript path from session_id if not provided + transcriptPath = i.computeTranscriptPath(ctx, sessionID, transcriptPath) + + return &agent.Event{ + Type: agent.TurnStart, + SessionID: sessionID, + SessionRef: transcriptPath, + Prompt: prompt, + Timestamp: time.Now(), + }, nil +} + +func (i *IFlowCLIAgent) parsePreToolUse(stdin io.Reader) (*agent.Event, error) { + var input ToolHookInput + if err := json.NewDecoder(stdin).Decode(&input); err != nil { + return nil, fmt.Errorf("failed to decode pre-tool-use input: %w", err) + } + + // Check if this is a subagent start (iFlow doesn't have explicit subagent concept, + // but we can detect certain patterns if needed) + // For now, we don't generate lifecycle events for PreToolUse + // unless it's a special tool that indicates subagent behavior + + return nil, nil +} + +func (i *IFlowCLIAgent) parsePostToolUse(stdin io.Reader) (*agent.Event, error) { + var input ToolHookInput + if err := json.NewDecoder(stdin).Decode(&input); err != nil { + return nil, fmt.Errorf("failed to decode post-tool-use input: %w", err) + } + + // Similar to PreToolUse, we don't generate lifecycle events for PostToolUse + // unless special handling is needed + + return nil, nil +} + +func (i *IFlowCLIAgent) parseStop(ctx context.Context, stdin io.Reader) (*agent.Event, error) { + var input StopInput + if err := json.NewDecoder(stdin).Decode(&input); err != nil { + return nil, fmt.Errorf("failed to decode stop input: %w", err) + } + + // Override with environment variables if present (iFlow CLI may use env vars) + sessionID := input.SessionID + if envSessionID := os.Getenv(EnvIFlowSessionID); envSessionID != "" { + sessionID = envSessionID + } + transcriptPath := input.TranscriptPath + if envTranscriptPath := os.Getenv(EnvIFlowTranscriptPath); envTranscriptPath != "" { + transcriptPath = envTranscriptPath + } + + // Fallback: compute transcript path from session_id if not provided + transcriptPath = i.computeTranscriptPath(ctx, sessionID, transcriptPath) + + event := &agent.Event{ + Type: agent.TurnEnd, + SessionID: sessionID, + SessionRef: transcriptPath, + Timestamp: time.Now(), + } + + if input.DurationMs > 0 { + event.DurationMs = input.DurationMs + } + if input.TurnCount > 0 { + event.TurnCount = input.TurnCount + } + + return event, nil +} + +func (i *IFlowCLIAgent) parseSessionEnd(ctx context.Context, stdin io.Reader) (*agent.Event, error) { + var input BaseHookInput + if err := json.NewDecoder(stdin).Decode(&input); err != nil { + return nil, fmt.Errorf("failed to decode session end input: %w", err) + } + + // Override with environment variables if present (iFlow CLI may use env vars) + sessionID := input.SessionID + if envSessionID := os.Getenv(EnvIFlowSessionID); envSessionID != "" { + sessionID = envSessionID + } + transcriptPath := input.TranscriptPath + if envTranscriptPath := os.Getenv(EnvIFlowTranscriptPath); envTranscriptPath != "" { + transcriptPath = envTranscriptPath + } + + // Fallback: compute transcript path from session_id if not provided + transcriptPath = i.computeTranscriptPath(ctx, sessionID, transcriptPath) + + return &agent.Event{ + Type: agent.SessionEnd, + SessionID: sessionID, + SessionRef: transcriptPath, + Timestamp: time.Now(), + }, nil +} + +func (i *IFlowCLIAgent) parseSubagentStop(ctx context.Context, stdin io.Reader) (*agent.Event, error) { + var input SubagentStopInput + if err := json.NewDecoder(stdin).Decode(&input); err != nil { + return nil, fmt.Errorf("failed to decode subagent stop input: %w", err) + } + + // Override with environment variables if present (iFlow CLI may use env vars) + sessionID := input.SessionID + if envSessionID := os.Getenv(EnvIFlowSessionID); envSessionID != "" { + sessionID = envSessionID + } + transcriptPath := input.TranscriptPath + if envTranscriptPath := os.Getenv(EnvIFlowTranscriptPath); envTranscriptPath != "" { + transcriptPath = envTranscriptPath + } + + // Fallback: compute transcript path from session_id if not provided + transcriptPath = i.computeTranscriptPath(ctx, sessionID, transcriptPath) + + event := &agent.Event{ + Type: agent.SubagentEnd, + SessionID: sessionID, + SessionRef: transcriptPath, + Timestamp: time.Now(), + } + + if input.SubagentID != "" { + event.SubagentID = input.SubagentID + } + if input.DurationMs > 0 { + event.DurationMs = input.DurationMs + } + + return event, nil +} diff --git a/cmd/entire/cli/agent/iflow/lifecycle_test.go b/cmd/entire/cli/agent/iflow/lifecycle_test.go new file mode 100644 index 000000000..ebcd12a02 --- /dev/null +++ b/cmd/entire/cli/agent/iflow/lifecycle_test.go @@ -0,0 +1,455 @@ +package iflow + +import ( + "context" + "encoding/json" + "strings" + "testing" + "time" + + "github.com/entireio/cli/cmd/entire/cli/agent" +) + +func TestParseSessionStart(t *testing.T) { + ag := NewIFlowCLIAgent().(*IFlowCLIAgent) + + tests := []struct { + name string + input map[string]interface{} + expected agent.EventType + wantErr bool + }{ + { + name: "valid session start", + input: map[string]interface{}{ + "session_id": "test-session", + "cwd": "/home/user/project", + "hook_event_name": "SessionStart", + "transcript_path": "/home/user/.iflow/projects/test-session.jsonl", + "source": "startup", + "model": "glm-5", + }, + expected: agent.SessionStart, + wantErr: false, + }, + { + name: "session start without model", + input: map[string]interface{}{ + "session_id": "test-session", + "cwd": "/home/user/project", + "hook_event_name": "SessionStart", + }, + expected: agent.SessionStart, + wantErr: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + data, _ := json.Marshal(tt.input) + stdin := strings.NewReader(string(data)) + + event, err := ag.ParseHookEvent(context.Background(), HookNameSessionStart, stdin) + if (err != nil) != tt.wantErr { + t.Errorf("ParseHookEvent() error = %v, wantErr %v", err, tt.wantErr) + return + } + if err != nil { + return + } + + if event.Type != tt.expected { + t.Errorf("Expected event type %v, got %v", tt.expected, event.Type) + } + + if event.SessionID != tt.input["session_id"] { + t.Errorf("Expected session ID %q, got %q", tt.input["session_id"], event.SessionID) + } + + if model, ok := tt.input["model"].(string); ok && event.Model != model { + t.Errorf("Expected model %q, got %q", model, event.Model) + } + }) + } +} + +func TestParseUserPromptSubmit(t *testing.T) { + ag := NewIFlowCLIAgent().(*IFlowCLIAgent) + + input := map[string]interface{}{ + "session_id": "test-session", + "cwd": "/home/user/project", + "hook_event_name": "UserPromptSubmit", + "transcript_path": "/home/user/.iflow/projects/test-session.jsonl", + "prompt": "Write a function to calculate fibonacci numbers", + } + + data, _ := json.Marshal(input) + stdin := strings.NewReader(string(data)) + + event, err := ag.ParseHookEvent(context.Background(), HookNameUserPromptSubmit, stdin) + if err != nil { + t.Fatalf("ParseHookEvent failed: %v", err) + } + + if event.Type != agent.TurnStart { + t.Errorf("Expected event type TurnStart, got %v", event.Type) + } + + if event.Prompt != input["prompt"] { + t.Errorf("Expected prompt %q, got %q", input["prompt"], event.Prompt) + } +} + +func TestParseStop(t *testing.T) { + ag := NewIFlowCLIAgent().(*IFlowCLIAgent) + + tests := []struct { + name string + input map[string]interface{} + expectedType agent.EventType + expectedDur int64 + expectedTurns int + }{ + { + name: "stop with metrics", + input: map[string]interface{}{ + "session_id": "test-session", + "cwd": "/home/user/project", + "hook_event_name": "Stop", + "transcript_path": "/home/user/.iflow/projects/test-session.jsonl", + "duration_ms": 5000, + "turn_count": 3, + }, + expectedType: agent.TurnEnd, + expectedDur: 5000, + expectedTurns: 3, + }, + { + name: "stop without metrics", + input: map[string]interface{}{ + "session_id": "test-session", + "cwd": "/home/user/project", + "hook_event_name": "Stop", + }, + expectedType: agent.TurnEnd, + expectedDur: 0, + expectedTurns: 0, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + data, _ := json.Marshal(tt.input) + stdin := strings.NewReader(string(data)) + + event, err := ag.ParseHookEvent(context.Background(), HookNameStop, stdin) + if err != nil { + t.Fatalf("ParseHookEvent failed: %v", err) + } + + if event.Type != tt.expectedType { + t.Errorf("Expected event type %v, got %v", tt.expectedType, event.Type) + } + + if event.DurationMs != tt.expectedDur { + t.Errorf("Expected duration %d, got %d", tt.expectedDur, event.DurationMs) + } + + if event.TurnCount != tt.expectedTurns { + t.Errorf("Expected turn count %d, got %d", tt.expectedTurns, event.TurnCount) + } + }) + } +} + +func TestParseSessionEnd(t *testing.T) { + ag := NewIFlowCLIAgent().(*IFlowCLIAgent) + + input := map[string]interface{}{ + "session_id": "test-session", + "cwd": "/home/user/project", + "hook_event_name": "SessionEnd", + "transcript_path": "/home/user/.iflow/projects/test-session.jsonl", + } + + data, _ := json.Marshal(input) + stdin := strings.NewReader(string(data)) + + event, err := ag.ParseHookEvent(context.Background(), HookNameSessionEnd, stdin) + if err != nil { + t.Fatalf("ParseHookEvent failed: %v", err) + } + + if event.Type != agent.SessionEnd { + t.Errorf("Expected event type SessionEnd, got %v", event.Type) + } + + if event.SessionID != input["session_id"] { + t.Errorf("Expected session ID %q, got %q", input["session_id"], event.SessionID) + } +} + +func TestParseSubagentStop(t *testing.T) { + ag := NewIFlowCLIAgent().(*IFlowCLIAgent) + + input := map[string]interface{}{ + "session_id": "test-session", + "cwd": "/home/user/project", + "hook_event_name": "SubagentStop", + "transcript_path": "/home/user/.iflow/projects/test-session.jsonl", + "subagent_id": "subagent-123", + "duration_ms": 2000, + } + + data, _ := json.Marshal(input) + stdin := strings.NewReader(string(data)) + + event, err := ag.ParseHookEvent(context.Background(), HookNameSubagentStop, stdin) + if err != nil { + t.Fatalf("ParseHookEvent failed: %v", err) + } + + if event.Type != agent.SubagentEnd { + t.Errorf("Expected event type SubagentEnd, got %v", event.Type) + } + + if event.SubagentID != input["subagent_id"] { + t.Errorf("Expected subagent ID %q, got %q", input["subagent_id"], event.SubagentID) + } + + if event.DurationMs != int64(input["duration_ms"].(int)) { + t.Errorf("Expected duration %d, got %d", input["duration_ms"], event.DurationMs) + } +} + +func TestParseHookEvent_UnknownHook(t *testing.T) { + ag := NewIFlowCLIAgent().(*IFlowCLIAgent) + + // Unknown hooks should return nil, nil + event, err := ag.ParseHookEvent(context.Background(), "unknown-hook", strings.NewReader("{}")) + if err != nil { + t.Errorf("Expected no error for unknown hook, got %v", err) + } + if event != nil { + t.Errorf("Expected nil event for unknown hook, got %v", event) + } +} + +func TestParseHookEvent_NonLifecycleHooks(t *testing.T) { + ag := NewIFlowCLIAgent().(*IFlowCLIAgent) + + // SetUpEnvironment and Notification should return nil, nil + nonLifecycleHooks := []string{HookNameSetUpEnvironment, HookNameNotification} + + for _, hookName := range nonLifecycleHooks { + t.Run(hookName, func(t *testing.T) { + event, err := ag.ParseHookEvent(context.Background(), hookName, strings.NewReader("{}")) + if err != nil { + t.Errorf("Expected no error for %s, got %v", hookName, err) + } + if event != nil { + t.Errorf("Expected nil event for %s, got %v", hookName, event) + } + }) + } +} + +func TestWriteHookResponse(t *testing.T) { + ag := NewIFlowCLIAgent().(*IFlowCLIAgent) + + // This test just verifies the function doesn't panic + // In a real test, we'd capture stdout + err := ag.WriteHookResponse("Test message") + if err != nil { + t.Errorf("WriteHookResponse failed: %v", err) + } +} + +func TestEventTimestamp(t *testing.T) { + ag := NewIFlowCLIAgent().(*IFlowCLIAgent) + + input := map[string]interface{}{ + "session_id": "test-session", + "cwd": "/home/user/project", + "hook_event_name": "SessionStart", + } + + data, _ := json.Marshal(input) + stdin := strings.NewReader(string(data)) + + before := time.Now() + event, err := ag.ParseHookEvent(context.Background(), HookNameSessionStart, stdin) + after := time.Now() + + if err != nil { + t.Fatalf("ParseHookEvent failed: %v", err) + } + + if event.Timestamp.Before(before) || event.Timestamp.After(after) { + t.Error("Event timestamp is not within expected range") + } +} + +func TestParsePreToolUse(t *testing.T) { + ag := NewIFlowCLIAgent().(*IFlowCLIAgent) + + tests := []struct { + name string + input map[string]interface{} + }{ + { + name: "pre-tool-use with file_path", + input: map[string]interface{}{ + "session_id": "test-session", + "cwd": "/home/user/project", + "hook_event_name": "PreToolUse", + "transcript_path": "/home/user/.iflow/projects/test-session.jsonl", + "tool_name": "write_file", + "tool_aliases": []string{"write", "create"}, + "tool_input": map[string]interface{}{ + "file_path": "/home/user/project/test.go", + "content": "package main", + }, + }, + }, + { + name: "pre-tool-use with edit tool", + input: map[string]interface{}{ + "session_id": "test-session", + "cwd": "/home/user/project", + "hook_event_name": "PreToolUse", + "tool_name": "replace", + "tool_aliases": []string{"Edit", "edit"}, + "tool_input": map[string]interface{}{ + "file_path": "/home/user/project/main.go", + "old_string": "old", + "new_string": "new", + }, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + data, _ := json.Marshal(tt.input) + stdin := strings.NewReader(string(data)) + + // PreToolUse currently returns nil (no lifecycle event) + event, err := ag.ParseHookEvent(context.Background(), HookNamePreToolUse, stdin) + if err != nil { + t.Errorf("ParseHookEvent failed: %v", err) + } + // PreToolUse doesn't generate a lifecycle event in current implementation + if event != nil { + t.Errorf("Expected nil event for PreToolUse, got %v", event) + } + }) + } +} + +func TestParsePostToolUse(t *testing.T) { + ag := NewIFlowCLIAgent().(*IFlowCLIAgent) + + tests := []struct { + name string + input map[string]interface{} + }{ + { + name: "post-tool-use with response", + input: map[string]interface{}{ + "session_id": "test-session", + "cwd": "/home/user/project", + "hook_event_name": "PostToolUse", + "transcript_path": "/home/user/.iflow/projects/test-session.jsonl", + "tool_name": "write_file", + "tool_aliases": []string{"write", "create"}, + "tool_input": map[string]interface{}{ + "file_path": "/home/user/project/test.go", + "content": "package main", + }, + "tool_response": map[string]interface{}{ + "result": map[string]interface{}{ + "llmContent": "File written successfully", + }, + }, + }, + }, + { + name: "post-tool-use without response", + input: map[string]interface{}{ + "session_id": "test-session", + "cwd": "/home/user/project", + "hook_event_name": "PostToolUse", + "tool_name": "read_file", + "tool_input": map[string]interface{}{ + "file_path": "/home/user/project/main.go", + }, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + data, _ := json.Marshal(tt.input) + stdin := strings.NewReader(string(data)) + + // PostToolUse currently returns nil (no lifecycle event) + event, err := ag.ParseHookEvent(context.Background(), HookNamePostToolUse, stdin) + if err != nil { + t.Errorf("ParseHookEvent failed: %v", err) + } + // PostToolUse doesn't generate a lifecycle event in current implementation + if event != nil { + t.Errorf("Expected nil event for PostToolUse, got %v", event) + } + }) + } +} + +func TestParseSetUpEnvironment(t *testing.T) { + ag := NewIFlowCLIAgent().(*IFlowCLIAgent) + + input := map[string]interface{}{ + "session_id": "test-session", + "cwd": "/home/user/project", + "hook_event_name": "SetUpEnvironment", + "transcript_path": "/home/user/.iflow/projects/test-session.jsonl", + } + + data, _ := json.Marshal(input) + stdin := strings.NewReader(string(data)) + + // SetUpEnvironment returns nil (no lifecycle event) + event, err := ag.ParseHookEvent(context.Background(), HookNameSetUpEnvironment, stdin) + if err != nil { + t.Errorf("ParseHookEvent failed: %v", err) + } + if event != nil { + t.Errorf("Expected nil event for SetUpEnvironment, got %v", event) + } +} + +func TestParseNotification(t *testing.T) { + ag := NewIFlowCLIAgent().(*IFlowCLIAgent) + + input := map[string]interface{}{ + "session_id": "test-session", + "cwd": "/home/user/project", + "hook_event_name": "Notification", + "transcript_path": "/home/user/.iflow/projects/test-session.jsonl", + "message": "Permission request for file access", + } + + data, _ := json.Marshal(input) + stdin := strings.NewReader(string(data)) + + // Notification returns nil (no lifecycle event) + event, err := ag.ParseHookEvent(context.Background(), HookNameNotification, stdin) + if err != nil { + t.Errorf("ParseHookEvent failed: %v", err) + } + if event != nil { + t.Errorf("Expected nil event for Notification, got %v", event) + } +} diff --git a/cmd/entire/cli/agent/iflow/transcript.go b/cmd/entire/cli/agent/iflow/transcript.go new file mode 100644 index 000000000..1a2474488 --- /dev/null +++ b/cmd/entire/cli/agent/iflow/transcript.go @@ -0,0 +1,235 @@ +package iflow + +import ( + "bufio" + "encoding/json" + "fmt" + "io" + "os" + "strings" +) + +// ParseTranscriptFromBytes parses JSONL transcript data into TranscriptLines. +func ParseTranscriptFromBytes(data []byte) ([]TranscriptLine, error) { + var lines []TranscriptLine + scanner := bufio.NewScanner(bytesReader(data)) + for scanner.Scan() { + line := strings.TrimSpace(scanner.Text()) + if line == "" { + continue + } + var entry TranscriptLine + if err := json.Unmarshal([]byte(line), &entry); err != nil { + // Skip malformed lines + continue + } + lines = append(lines, entry) + } + if err := scanner.Err(); err != nil { + return nil, fmt.Errorf("failed to scan transcript: %w", err) + } + return lines, nil +} + +// ExtractModifiedFiles extracts file paths from tools that modify files. +func ExtractModifiedFiles(lines []TranscriptLine) []string { + fileSet := make(map[string]struct{}) + + for _, line := range lines { + if line.ToolUse == nil { + continue + } + + if !isFileModificationTool(line.ToolUse.Name) { + continue + } + + filePath := extractFilePathFromToolInput(line.ToolUse.Input) + if filePath != "" { + fileSet[filePath] = struct{}{} + } + } + + // Convert set to sorted slice + files := make([]string, 0, len(fileSet)) + for file := range fileSet { + files = append(files, file) + } + return files +} + +// isFileModificationTool checks if the tool name indicates a file modification operation. +func isFileModificationTool(toolName string) bool { + for _, t := range FileModificationTools { + if toolName == t { + return true + } + } + return false +} + +// extractFilePathFromToolInput extracts the file path from tool input JSON. +func extractFilePathFromToolInput(input json.RawMessage) string { + if len(input) == 0 { + return "" + } + + // Try FileWriteToolInput structure + var writeInput FileWriteToolInput + if err := json.Unmarshal(input, &writeInput); err == nil { + if writeInput.FilePath != "" { + return writeInput.FilePath + } + if writeInput.Path != "" { + return writeInput.Path + } + } + + // Try FileEditToolInput structure + var editInput FileEditToolInput + if err := json.Unmarshal(input, &editInput); err == nil { + if editInput.FilePath != "" { + return editInput.FilePath + } + if editInput.Path != "" { + return editInput.Path + } + } + + // Try generic map extraction + var generic map[string]interface{} + if err := json.Unmarshal(input, &generic); err == nil { + // Check common field names + for _, key := range []string{"file_path", "path", "filepath", "file", "filename"} { + if val, ok := generic[key]; ok { + if str, ok := val.(string); ok && str != "" { + return str + } + } + } + } + + return "" +} + +// bytesReader creates an io.Reader from bytes (helper for testing). +func bytesReader(data []byte) io.Reader { + return strings.NewReader(string(data)) +} + +// SerializeTranscript converts TranscriptLines back to JSONL format. +func SerializeTranscript(lines []TranscriptLine) ([]byte, error) { + var result strings.Builder + encoder := json.NewEncoder(&result) + encoder.SetEscapeHTML(false) + + for _, line := range lines { + if err := encoder.Encode(line); err != nil { + return nil, fmt.Errorf("failed to encode transcript line: %w", err) + } + } + + return []byte(result.String()), nil +} + +// FindCheckpointLine finds the line index containing a specific tool result. +// Used for checkpoint rewind operations. +func FindCheckpointLine(lines []TranscriptLine, toolUseID string) (int, bool) { + for i, line := range lines { + if line.ToolResult != nil && line.ToolResult.ToolUseID == toolUseID { + return i, true + } + } + return 0, false +} + +// TruncateAtLine returns a new transcript truncated at the given line index (inclusive). +func TruncateAtLine(lines []TranscriptLine, lineIndex int) []TranscriptLine { + if lineIndex < 0 || lineIndex >= len(lines) { + return lines + } + return lines[:lineIndex+1] +} + +// TruncateAtToolResult returns a new transcript truncated at the given tool result. +func TruncateAtToolResult(lines []TranscriptLine, toolUseID string) ([]TranscriptLine, bool) { + lineIndex, found := FindCheckpointLine(lines, toolUseID) + if !found { + return nil, false + } + return TruncateAtLine(lines, lineIndex), true +} + +// GetLastToolUse finds the last tool use in the transcript. +func GetLastToolUse(lines []TranscriptLine) *ToolUse { + for i := len(lines) - 1; i >= 0; i-- { + if lines[i].ToolUse != nil { + return lines[i].ToolUse + } + } + return nil +} + +// GetToolResultForUse finds the tool result for a given tool use ID. +func GetToolResultForUse(lines []TranscriptLine, toolUseID string) *ToolResult { + for _, line := range lines { + if line.ToolResult != nil && line.ToolResult.ToolUseID == toolUseID { + return line.ToolResult + } + } + return nil +} + +// CountToolUses counts the occurrences of a specific tool in the transcript. +func CountToolUses(lines []TranscriptLine, toolName string) int { + count := 0 + for _, line := range lines { + if line.ToolUse != nil && line.ToolUse.Name == toolName { + count++ + } + } + return count +} + +// ReadTranscriptFile reads and parses a transcript file. +func ReadTranscriptFile(path string) ([]TranscriptLine, error) { + data, err := os.ReadFile(path) + if err != nil { + return nil, fmt.Errorf("failed to read transcript file: %w", err) + } + return ParseTranscriptFromBytes(data) +} + +// WriteTranscriptFile writes transcript lines to a file. +func WriteTranscriptFile(path string, lines []TranscriptLine) error { + data, err := SerializeTranscript(lines) + if err != nil { + return fmt.Errorf("failed to serialize transcript: %w", err) + } + + if err := os.WriteFile(path, data, 0o600); err != nil { + return fmt.Errorf("failed to write transcript file: %w", err) + } + + return nil +} + +// MergeTranscripts merges multiple transcripts into one. +// Handles deduplication of lines based on timestamp and content. +func MergeTranscripts(transcripts [][]TranscriptLine) []TranscriptLine { + seen := make(map[string]struct{}) + var result []TranscriptLine + + for _, transcript := range transcripts { + for _, line := range transcript { + // Create a key for deduplication + key := line.Timestamp + string(line.Message) + if _, exists := seen[key]; !exists { + seen[key] = struct{}{} + result = append(result, line) + } + } + } + + return result +} diff --git a/cmd/entire/cli/agent/iflow/transcript_test.go b/cmd/entire/cli/agent/iflow/transcript_test.go new file mode 100644 index 000000000..5bede61cf --- /dev/null +++ b/cmd/entire/cli/agent/iflow/transcript_test.go @@ -0,0 +1,441 @@ +package iflow + +import ( + "encoding/json" + "os" + "path/filepath" + "testing" +) + +func TestParseTranscriptFromBytes(t *testing.T) { + transcript := `{"type": "user", "timestamp": "2024-01-01T00:00:00Z", "message": {}} +{"type": "assistant", "timestamp": "2024-01-01T00:00:01Z", "message": {}} +{"type": "tool_use", "timestamp": "2024-01-01T00:00:02Z", "tool_use": {"id": "tool-1", "name": "write_file", "input": {"file_path": "test.txt"}}}` + + lines, err := ParseTranscriptFromBytes([]byte(transcript)) + if err != nil { + t.Fatalf("ParseTranscriptFromBytes failed: %v", err) + } + + if len(lines) != 3 { + t.Errorf("Expected 3 lines, got %d", len(lines)) + } + + if lines[2].ToolUse == nil { + t.Fatal("Expected tool use on line 3") + } + + if lines[2].ToolUse.Name != "write_file" { + t.Errorf("Expected tool name 'write_file', got %q", lines[2].ToolUse.Name) + } +} + +func TestParseTranscriptFromBytes_Empty(t *testing.T) { + lines, err := ParseTranscriptFromBytes([]byte("")) + if err != nil { + t.Fatalf("ParseTranscriptFromBytes failed: %v", err) + } + + if len(lines) != 0 { + t.Errorf("Expected 0 lines for empty input, got %d", len(lines)) + } +} + +func TestParseTranscriptFromBytes_InvalidJSON(t *testing.T) { + // Invalid JSON lines should be skipped + transcript := `{"type": "user", "timestamp": "2024-01-01T00:00:00Z"} +invalid json line +{"type": "assistant", "timestamp": "2024-01-01T00:00:01Z"}` + + lines, err := ParseTranscriptFromBytes([]byte(transcript)) + if err != nil { + t.Fatalf("ParseTranscriptFromBytes failed: %v", err) + } + + if len(lines) != 2 { + t.Errorf("Expected 2 valid lines, got %d", len(lines)) + } +} + +func TestExtractModifiedFiles(t *testing.T) { + lines := []TranscriptLine{ + { + Type: "tool_use", + Timestamp: "2024-01-01T00:00:00Z", + ToolUse: &ToolUse{ + ID: "tool-1", + Name: "write_file", + Input: json.RawMessage(`{"file_path": "file1.txt"}`), + }, + }, + { + Type: "tool_use", + Timestamp: "2024-01-01T00:00:01Z", + ToolUse: &ToolUse{ + ID: "tool-2", + Name: "replace", + Input: json.RawMessage(`{"path": "file2.txt"}`), + }, + }, + { + Type: "tool_use", + Timestamp: "2024-01-01T00:00:02Z", + ToolUse: &ToolUse{ + ID: "tool-3", + Name: "read_file", + Input: json.RawMessage(`{"file_path": "file3.txt"}`), + }, + }, + } + + files := ExtractModifiedFiles(lines) + + if len(files) != 2 { + t.Errorf("Expected 2 modified files, got %d", len(files)) + } + + // Check that file1.txt and file2.txt are in the list + found := make(map[string]bool) + for _, f := range files { + found[f] = true + } + + if !found["file1.txt"] { + t.Error("Expected file1.txt in modified files") + } + if !found["file2.txt"] { + t.Error("Expected file2.txt in modified files") + } + if found["file3.txt"] { + t.Error("file3.txt should not be in modified files (read_file is not a modification tool)") + } +} + +func TestExtractModifiedFiles_Duplicates(t *testing.T) { + lines := []TranscriptLine{ + { + Type: "tool_use", + Timestamp: "2024-01-01T00:00:00Z", + ToolUse: &ToolUse{ + ID: "tool-1", + Name: "write_file", + Input: json.RawMessage(`{"file_path": "same.txt"}`), + }, + }, + { + Type: "tool_use", + Timestamp: "2024-01-01T00:00:01Z", + ToolUse: &ToolUse{ + ID: "tool-2", + Name: "write_file", + Input: json.RawMessage(`{"file_path": "same.txt"}`), + }, + }, + } + + files := ExtractModifiedFiles(lines) + + if len(files) != 1 { + t.Errorf("Expected 1 unique file, got %d", len(files)) + } + + if files[0] != "same.txt" { + t.Errorf("Expected 'same.txt', got %q", files[0]) + } +} + +func TestIsFileModificationTool(t *testing.T) { + tests := []struct { + toolName string + expected bool + }{ + {"write_file", true}, + {"replace", true}, + {"multi_edit", true}, + {"read_file", false}, + {"run_shell_command", false}, + {"list_directory", false}, + } + + for _, tt := range tests { + t.Run(tt.toolName, func(t *testing.T) { + result := isFileModificationTool(tt.toolName) + if result != tt.expected { + t.Errorf("isFileModificationTool(%q) = %v, want %v", tt.toolName, result, tt.expected) + } + }) + } +} + +func TestExtractFilePathFromToolInput(t *testing.T) { + tests := []struct { + name string + input map[string]interface{} + expected string + }{ + { + name: "file_path field", + input: map[string]interface{}{"file_path": "test.txt"}, + expected: "test.txt", + }, + { + name: "path field", + input: map[string]interface{}{"path": "src/main.go"}, + expected: "src/main.go", + }, + { + name: "filepath field", + input: map[string]interface{}{"filepath": "docs/readme.md"}, + expected: "docs/readme.md", + }, + { + name: "no file path", + input: map[string]interface{}{"content": "hello"}, + expected: "", + }, + { + name: "empty input", + input: map[string]interface{}{}, + expected: "", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + data, _ := json.Marshal(tt.input) + result := extractFilePathFromToolInput(data) + if result != tt.expected { + t.Errorf("extractFilePathFromToolInput() = %q, want %q", result, tt.expected) + } + }) + } +} + +func TestSerializeTranscript(t *testing.T) { + lines := []TranscriptLine{ + { + Type: "user", + Timestamp: "2024-01-01T00:00:00Z", + }, + { + Type: "assistant", + Timestamp: "2024-01-01T00:00:01Z", + }, + } + + data, err := SerializeTranscript(lines) + if err != nil { + t.Fatalf("SerializeTranscript failed: %v", err) + } + + // Parse it back to verify + parsed, err := ParseTranscriptFromBytes(data) + if err != nil { + t.Fatalf("ParseTranscriptFromBytes failed: %v", err) + } + + if len(parsed) != 2 { + t.Errorf("Expected 2 lines after round-trip, got %d", len(parsed)) + } +} + +func TestFindCheckpointLine(t *testing.T) { + lines := []TranscriptLine{ + {Type: "user", Timestamp: "2024-01-01T00:00:00Z"}, + { + Type: "tool_result", + Timestamp: "2024-01-01T00:00:01Z", + ToolResult: &ToolResult{ToolUseID: "tool-1"}, + }, + {Type: "assistant", Timestamp: "2024-01-01T00:00:02Z"}, + { + Type: "tool_result", + Timestamp: "2024-01-01T00:00:03Z", + ToolResult: &ToolResult{ToolUseID: "tool-2"}, + }, + } + + index, found := FindCheckpointLine(lines, "tool-2") + if !found { + t.Error("Expected to find tool-2") + } + if index != 3 { + t.Errorf("Expected index 3, got %d", index) + } + + _, found = FindCheckpointLine(lines, "non-existent") + if found { + t.Error("Expected not to find non-existent tool") + } +} + +func TestTruncateAtLine(t *testing.T) { + lines := []TranscriptLine{ + {Type: "user", Timestamp: "2024-01-01T00:00:00Z"}, + {Type: "assistant", Timestamp: "2024-01-01T00:00:01Z"}, + {Type: "user", Timestamp: "2024-01-01T00:00:02Z"}, + } + + truncated := TruncateAtLine(lines, 1) + if len(truncated) != 2 { + t.Errorf("Expected 2 lines after truncate, got %d", len(truncated)) + } + + // Test out of bounds + truncated = TruncateAtLine(lines, 10) + if len(truncated) != 3 { + t.Errorf("Expected original lines for out of bounds, got %d", len(truncated)) + } +} + +func TestReadTranscriptFile(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "test.jsonl") + + content := `{"type": "user", "timestamp": "2024-01-01T00:00:00Z"} +{"type": "assistant", "timestamp": "2024-01-01T00:00:01Z"}` + + if err := os.WriteFile(path, []byte(content), 0o600); err != nil { + t.Fatal(err) + } + + lines, err := ReadTranscriptFile(path) + if err != nil { + t.Fatalf("ReadTranscriptFile failed: %v", err) + } + + if len(lines) != 2 { + t.Errorf("Expected 2 lines, got %d", len(lines)) + } +} + +func TestWriteTranscriptFile(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "test.jsonl") + + lines := []TranscriptLine{ + {Type: "user", Timestamp: "2024-01-01T00:00:00Z"}, + {Type: "assistant", Timestamp: "2024-01-01T00:00:01Z"}, + } + + if err := WriteTranscriptFile(path, lines); err != nil { + t.Fatalf("WriteTranscriptFile failed: %v", err) + } + + // Read it back + data, err := os.ReadFile(path) + if err != nil { + t.Fatal(err) + } + + parsed, err := ParseTranscriptFromBytes(data) + if err != nil { + t.Fatalf("ParseTranscriptFromBytes failed: %v", err) + } + + if len(parsed) != 2 { + t.Errorf("Expected 2 lines after write/read, got %d", len(parsed)) + } +} + +func TestMergeTranscripts(t *testing.T) { + transcript1 := []TranscriptLine{ + {Type: "user", Timestamp: "2024-01-01T00:00:00Z"}, + {Type: "assistant", Timestamp: "2024-01-01T00:00:01Z"}, + } + + transcript2 := []TranscriptLine{ + {Type: "assistant", Timestamp: "2024-01-01T00:00:01Z"}, // Duplicate + {Type: "user", Timestamp: "2024-01-01T00:00:02Z"}, + } + + merged := MergeTranscripts([][]TranscriptLine{transcript1, transcript2}) + + if len(merged) != 3 { + t.Errorf("Expected 3 unique lines after merge, got %d", len(merged)) + } +} + +func TestGetLastToolUse(t *testing.T) { + lines := []TranscriptLine{ + {Type: "user", Timestamp: "2024-01-01T00:00:00Z"}, + { + Type: "tool_use", + Timestamp: "2024-01-01T00:00:01Z", + ToolUse: &ToolUse{ID: "tool-1", Name: "write_file"}, + }, + {Type: "assistant", Timestamp: "2024-01-01T00:00:02Z"}, + } + + toolUse := GetLastToolUse(lines) + if toolUse == nil { + t.Fatal("Expected to find last tool use") + } + if toolUse.ID != "tool-1" { + t.Errorf("Expected tool-1, got %q", toolUse.ID) + } +} + +func TestGetToolResultForUse(t *testing.T) { + lines := []TranscriptLine{ + { + Type: "tool_use", + Timestamp: "2024-01-01T00:00:00Z", + ToolUse: &ToolUse{ID: "tool-1", Name: "write_file"}, + }, + { + Type: "tool_result", + Timestamp: "2024-01-01T00:00:01Z", + ToolResult: &ToolResult{ToolUseID: "tool-1"}, + }, + } + + result := GetToolResultForUse(lines, "tool-1") + if result == nil { + t.Fatal("Expected to find tool result") + } + if result.ToolUseID != "tool-1" { + t.Errorf("Expected tool_use_id tool-1, got %q", result.ToolUseID) + } + + result = GetToolResultForUse(lines, "non-existent") + if result != nil { + t.Error("Expected nil for non-existent tool use") + } +} + +func TestCountToolUses(t *testing.T) { + lines := []TranscriptLine{ + {Type: "user", Timestamp: "2024-01-01T00:00:00Z"}, + { + Type: "tool_use", + Timestamp: "2024-01-01T00:00:01Z", + ToolUse: &ToolUse{ID: "tool-1", Name: "write_file"}, + }, + { + Type: "tool_use", + Timestamp: "2024-01-01T00:00:02Z", + ToolUse: &ToolUse{ID: "tool-2", Name: "write_file"}, + }, + { + Type: "tool_use", + Timestamp: "2024-01-01T00:00:03Z", + ToolUse: &ToolUse{ID: "tool-3", Name: "read_file"}, + }, + } + + count := CountToolUses(lines, "write_file") + if count != 2 { + t.Errorf("Expected 2 write_file uses, got %d", count) + } + + count = CountToolUses(lines, "read_file") + if count != 1 { + t.Errorf("Expected 1 read_file use, got %d", count) + } + + count = CountToolUses(lines, "non-existent") + if count != 0 { + t.Errorf("Expected 0 uses for non-existent tool, got %d", count) + } +} diff --git a/cmd/entire/cli/agent/iflow/types.go b/cmd/entire/cli/agent/iflow/types.go new file mode 100644 index 000000000..19839e754 --- /dev/null +++ b/cmd/entire/cli/agent/iflow/types.go @@ -0,0 +1,205 @@ +package iflow + +import "encoding/json" + +// IFlowSettings represents the .iflow/settings.json structure +type IFlowSettings struct { + Hooks IFlowHooks `json:"hooks,omitempty"` + Permissions IFlowPermissions `json:"permissions,omitempty"` +} + +// IFlowHooks contains the hook configurations +// All hooks use IFlowHookMatcher format (with hooks array wrapper) per iFlow CLI specification +type IFlowHooks struct { + PreToolUse []IFlowHookMatcher `json:"PreToolUse,omitempty"` + PostToolUse []IFlowHookMatcher `json:"PostToolUse,omitempty"` + SetUpEnvironment []IFlowHookMatcher `json:"SetUpEnvironment,omitempty"` + Stop []IFlowHookMatcher `json:"Stop,omitempty"` + SubagentStop []IFlowHookMatcher `json:"SubagentStop,omitempty"` + SessionStart []IFlowHookMatcher `json:"SessionStart,omitempty"` + SessionEnd []IFlowHookMatcher `json:"SessionEnd,omitempty"` + UserPromptSubmit []IFlowHookMatcher `json:"UserPromptSubmit,omitempty"` + Notification []IFlowHookMatcher `json:"Notification,omitempty"` +} + +// IFlowPermissions contains permission settings +type IFlowPermissions struct { + Deny []string `json:"deny,omitempty"` +} + +// IFlowHookMatcher matches hooks to specific patterns +type IFlowHookMatcher struct { + Matcher string `json:"matcher,omitempty"` + Hooks []IFlowHookEntry `json:"hooks"` +} + +// IFlowHookEntry represents a single hook command +type IFlowHookEntry struct { + Type string `json:"type"` + Command string `json:"command"` + Timeout int `json:"timeout,omitempty"` +} + +// --- Hook Input Types --- + +// BaseHookInput contains fields common to all hook inputs +type BaseHookInput struct { + SessionID string `json:"session_id"` + CWD string `json:"cwd"` + HookEventName string `json:"hook_event_name"` + TranscriptPath string `json:"transcript_path,omitempty"` +} + +// ToolHookInput is the JSON structure from PreToolUse/PostToolUse hooks +type ToolHookInput struct { + BaseHookInput + + ToolName string `json:"tool_name"` + ToolAliases []string `json:"tool_aliases,omitempty"` + ToolInput json.RawMessage `json:"tool_input"` + ToolResponse json.RawMessage `json:"tool_response,omitempty"` +} + +// ToolResponseContent represents the structure of tool_response.result +type ToolResponseContent struct { + LLMContent string `json:"llmContent,omitempty"` +} + +// UserPromptSubmitInput is the JSON structure from UserPromptSubmit hook +type UserPromptSubmitInput struct { + BaseHookInput + + Prompt string `json:"prompt"` +} + +// SessionStartInput is the JSON structure from SessionStart hook +type SessionStartInput struct { + BaseHookInput + + Source string `json:"source,omitempty"` // startup, resume, clear, compress + Model string `json:"model,omitempty"` +} + +// NotificationInput is the JSON structure from Notification hook +type NotificationInput struct { + BaseHookInput + + Message string `json:"message"` +} + +// StopInput is the JSON structure from Stop hook +type StopInput struct { + BaseHookInput + + DurationMs int64 `json:"duration_ms,omitempty"` + TurnCount int `json:"turn_count,omitempty"` +} + +// SubagentStopInput is the JSON structure from SubagentStop hook +type SubagentStopInput struct { + BaseHookInput + + SubagentID string `json:"subagent_id,omitempty"` + DurationMs int64 `json:"duration_ms,omitempty"` +} + +// --- Transcript Types --- + +// TranscriptLine represents a single line in the JSONL transcript +type TranscriptLine struct { + Type string `json:"type"` + Timestamp string `json:"timestamp"` + Message json.RawMessage `json:"message,omitempty"` + ToolUse *ToolUse `json:"tool_use,omitempty"` + ToolResult *ToolResult `json:"tool_result,omitempty"` + Tokens *IFlowMessageTokens `json:"tokens,omitempty"` +} + +// IFlowMessageTokens represents token usage from an iFlow API response. +// This structure may be present in transcript lines for assistant messages. +type IFlowMessageTokens struct { + Input int `json:"input,omitempty"` + Output int `json:"output,omitempty"` + Cached int `json:"cached,omitempty"` + Total int `json:"total,omitempty"` +} + +// ToolUse represents a tool invocation in the transcript +type ToolUse struct { + ID string `json:"id"` + Name string `json:"name"` + Input json.RawMessage `json:"input"` +} + +// ToolResult represents the result of a tool invocation +type ToolResult struct { + ToolUseID string `json:"tool_use_id"` + Result json.RawMessage `json:"result"` +} + +// FileEditToolInput represents input for file editing tools +type FileEditToolInput struct { + FilePath string `json:"file_path"` + Path string `json:"path"` // Alternative field name +} + +// FileWriteToolInput represents input for file write tools +type FileWriteToolInput struct { + FilePath string `json:"file_path"` + Path string `json:"path"` // Alternative field name +} + +// Tool names used in iFlow CLI transcripts +const ( + ToolWrite = "write_file" + ToolEdit = "replace" + ToolMultiEdit = "multi_edit" + ToolShell = "run_shell_command" + ToolRead = "read_file" + ToolList = "list_directory" + ToolSearch = "search_file_content" + ToolGlob = "glob" +) + +// FileModificationTools lists tools that create or modify files +var FileModificationTools = []string{ + ToolWrite, + ToolEdit, + ToolMultiEdit, +} + +// SessionSource constants for SessionStart hook +type SessionSource string + +const ( + SessionSourceStartup SessionSource = "startup" + SessionSourceResume SessionSource = "resume" + SessionSourceClear SessionSource = "clear" + SessionSourceCompress SessionSource = "compress" +) + +// IFlow-specific environment variable names +const ( + EnvIFlowSessionID = "IFLOW_SESSION_ID" + EnvIFlowTranscriptPath = "IFLOW_TRANSCRIPT_PATH" + EnvIFlowCWD = "IFLOW_CWD" + EnvIFlowHookEventName = "IFLOW_HOOK_EVENT_NAME" + EnvIFlowToolName = "IFLOW_TOOL_NAME" + EnvIFlowToolArgs = "IFLOW_TOOL_ARGS" + EnvIFlowToolAliases = "IFLOW_TOOL_ALIASES" + EnvIFlowSessionSource = "IFLOW_SESSION_SOURCE" + EnvIFlowUserPrompt = "IFLOW_USER_PROMPT" + EnvIFlowNotification = "IFLOW_NOTIFICATION_MESSAGE" +) + +// Settings file name +const IFlowSettingsFileName = "settings.json" + +// metadataDenyRule blocks iFlow from reading Entire session metadata +const metadataDenyRule = "Read(./.entire/metadata/**)" + +// entireHookPrefixes are command prefixes that identify Entire hooks +var entireHookPrefixes = []string{ + "entire ", + "go run ${IFLOW_PROJECT_DIR}/cmd/entire/main.go ", +} diff --git a/cmd/entire/cli/agent/registry.go b/cmd/entire/cli/agent/registry.go index 5d518470f..e0721287a 100644 --- a/cmd/entire/cli/agent/registry.go +++ b/cmd/entire/cli/agent/registry.go @@ -105,7 +105,9 @@ const ( AgentNameCursor types.AgentName = "cursor" AgentNameFactoryAIDroid types.AgentName = "factoryai-droid" AgentNameGemini types.AgentName = "gemini" + AgentNameIFlow types.AgentName = "iflow" AgentNameOpenCode types.AgentName = "opencode" + AgentNameTraeAgent types.AgentName = "trae-agent" ) // Agent type constants (type identifiers stored in metadata/trailers) @@ -115,7 +117,9 @@ const ( AgentTypeCursor types.AgentType = "Cursor" AgentTypeFactoryAIDroid types.AgentType = "Factory AI Droid" AgentTypeGemini types.AgentType = "Gemini CLI" + AgentTypeIFlow types.AgentType = "iFlow CLI" AgentTypeOpenCode types.AgentType = "OpenCode" + AgentTypeTraeAgent types.AgentType = "Trae Agent" AgentTypeUnknown types.AgentType = "Agent" // Fallback for backwards compatibility ) diff --git a/cmd/entire/cli/agent/traeagent/hooks.go b/cmd/entire/cli/agent/traeagent/hooks.go new file mode 100644 index 000000000..1b7075d04 --- /dev/null +++ b/cmd/entire/cli/agent/traeagent/hooks.go @@ -0,0 +1,568 @@ +// Package traeagent implements the Agent interface for Trae Agent. +package traeagent + +import ( + "context" + "encoding/json" + "fmt" + "io" + "os" + "path/filepath" + "strings" + "time" + + "github.com/entireio/cli/cmd/entire/cli/agent" + "github.com/entireio/cli/cmd/entire/cli/jsonutil" + "github.com/entireio/cli/cmd/entire/cli/paths" +) + +// Ensure TraeAgent implements HookSupport +var _ agent.HookSupport = (*TraeAgent)(nil) + +// Trae Agent hook names - these become subcommands under `entire hooks trae-agent` +const ( + HookNameSessionStart = "session-start" + HookNameSessionEnd = "session-end" + HookNameBeforeAgent = "before-agent" + HookNameAfterAgent = "after-agent" + HookNameBeforeModel = "before-model" + HookNameAfterModel = "after-model" + HookNameBeforeToolSelection = "before-tool-selection" + HookNamePreTool = "pre-tool" + HookNameAfterTool = "after-tool" + HookNamePreCompress = "pre-compress" + HookNameNotification = "notification" +) + +// TraeSettingsFileName is the settings file used by Trae Agent. +const TraeSettingsFileName = "settings.json" + +// GetHookNames returns the hook verbs Trae Agent supports. +// These become subcommands: entire hooks trae-agent +func (t *TraeAgent) HookNames() []string { + return []string{ + HookNameSessionStart, + HookNameSessionEnd, + HookNameBeforeAgent, + HookNameAfterAgent, + HookNameBeforeModel, + HookNameAfterModel, + HookNameBeforeToolSelection, + HookNamePreTool, + HookNameAfterTool, + HookNamePreCompress, + HookNameNotification, + } +} + +// entireHookPrefixes are command prefixes that identify Entire hooks (both old and new formats) +var entireHookPrefixes = []string{ + "entire ", + "go run ${TRAE_PROJECT_DIR}/cmd/entire/main.go ", +} + +// InstallHooks installs Trae Agent hooks in .trae/settings.json. +// If force is true, removes existing Entire hooks before installing. +// Returns the number of hooks installed. +func (t *TraeAgent) InstallHooks(ctx context.Context, localDev bool, force bool) (int, error) { + // Use repo root instead of CWD to find .trae directory + repoRoot, err := paths.WorktreeRoot(ctx) + if err != nil { + // Fallback to CWD if not in a git repo (e.g., during tests) + repoRoot, err = os.Getwd() //nolint:forbidigo // Intentional fallback when WorktreeRoot() fails + if err != nil { + return 0, fmt.Errorf("failed to get current directory: %w", err) + } + } + + settingsPath := filepath.Join(repoRoot, ".trae", TraeSettingsFileName) + + // Read existing settings if they exist + var rawSettings map[string]json.RawMessage + + // rawHooks preserves unknown hook types + var rawHooks map[string]json.RawMessage + + existingData, readErr := os.ReadFile(settingsPath) //nolint:gosec // path is constructed safely + if readErr == nil { + if err := json.Unmarshal(existingData, &rawSettings); err != nil { + return 0, fmt.Errorf("failed to parse existing settings.json: %w", err) + } + if hooksRaw, ok := rawSettings["hooks"]; ok { + if err := json.Unmarshal(hooksRaw, &rawHooks); err != nil { + return 0, fmt.Errorf("failed to parse hooks in settings.json: %w", err) + } + } + } else { + rawSettings = make(map[string]json.RawMessage) + } + + if rawHooks == nil { + rawHooks = make(map[string]json.RawMessage) + } + + // Parse only the hook types we need to modify + var sessionStart, sessionEnd, beforeAgent, afterAgent, beforeModel, afterModel []TraeHook + var beforeToolSelection, preTool, afterTool, preCompress, notification []TraeHook + + parseHookType(rawHooks, "SessionStart", &sessionStart) + parseHookType(rawHooks, "SessionEnd", &sessionEnd) + parseHookType(rawHooks, "BeforeAgent", &beforeAgent) + parseHookType(rawHooks, "AfterAgent", &afterAgent) + parseHookType(rawHooks, "BeforeModel", &beforeModel) + parseHookType(rawHooks, "AfterModel", &afterModel) + parseHookType(rawHooks, "BeforeToolSelection", &beforeToolSelection) + parseHookType(rawHooks, "PreTool", &preTool) + parseHookType(rawHooks, "AfterTool", &afterTool) + parseHookType(rawHooks, "PreCompress", &preCompress) + parseHookType(rawHooks, "Notification", ¬ification) + + // If force is true, remove all existing Entire hooks first + if force { + sessionStart = removeEntireHooks(sessionStart) + sessionEnd = removeEntireHooks(sessionEnd) + beforeAgent = removeEntireHooks(beforeAgent) + afterAgent = removeEntireHooks(afterAgent) + beforeModel = removeEntireHooks(beforeModel) + afterModel = removeEntireHooks(afterModel) + beforeToolSelection = removeEntireHooks(beforeToolSelection) + preTool = removeEntireHooks(preTool) + afterTool = removeEntireHooks(afterTool) + preCompress = removeEntireHooks(preCompress) + notification = removeEntireHooks(notification) + } + + // Define hook commands + var sessionStartCmd, sessionEndCmd, beforeAgentCmd, afterAgentCmd string + var beforeModelCmd, afterModelCmd, beforeToolSelectionCmd, preToolCmd string + var afterToolCmd, preCompressCmd, notificationCmd string + + if localDev { + baseCmd := "go run ${TRAE_PROJECT_DIR}/cmd/entire/main.go hooks trae-agent" + sessionStartCmd = fmt.Sprintf("%s %s", baseCmd, HookNameSessionStart) + sessionEndCmd = fmt.Sprintf("%s %s", baseCmd, HookNameSessionEnd) + beforeAgentCmd = fmt.Sprintf("%s %s", baseCmd, HookNameBeforeAgent) + afterAgentCmd = fmt.Sprintf("%s %s", baseCmd, HookNameAfterAgent) + beforeModelCmd = fmt.Sprintf("%s %s", baseCmd, HookNameBeforeModel) + afterModelCmd = fmt.Sprintf("%s %s", baseCmd, HookNameAfterModel) + beforeToolSelectionCmd = fmt.Sprintf("%s %s", baseCmd, HookNameBeforeToolSelection) + preToolCmd = fmt.Sprintf("%s %s", baseCmd, HookNamePreTool) + afterToolCmd = fmt.Sprintf("%s %s", baseCmd, HookNameAfterTool) + preCompressCmd = fmt.Sprintf("%s %s", baseCmd, HookNamePreCompress) + notificationCmd = fmt.Sprintf("%s %s", baseCmd, HookNameNotification) + } else { + baseCmd := "entire hooks trae-agent" + sessionStartCmd = fmt.Sprintf("%s %s", baseCmd, HookNameSessionStart) + sessionEndCmd = fmt.Sprintf("%s %s", baseCmd, HookNameSessionEnd) + beforeAgentCmd = fmt.Sprintf("%s %s", baseCmd, HookNameBeforeAgent) + afterAgentCmd = fmt.Sprintf("%s %s", baseCmd, HookNameAfterAgent) + beforeModelCmd = fmt.Sprintf("%s %s", baseCmd, HookNameBeforeModel) + afterModelCmd = fmt.Sprintf("%s %s", baseCmd, HookNameAfterModel) + beforeToolSelectionCmd = fmt.Sprintf("%s %s", baseCmd, HookNameBeforeToolSelection) + preToolCmd = fmt.Sprintf("%s %s", baseCmd, HookNamePreTool) + afterToolCmd = fmt.Sprintf("%s %s", baseCmd, HookNameAfterTool) + preCompressCmd = fmt.Sprintf("%s %s", baseCmd, HookNamePreCompress) + notificationCmd = fmt.Sprintf("%s %s", baseCmd, HookNameNotification) + } + + count := 0 + + // Add hooks if they don't exist + if !hookCommandExists(sessionStart, sessionStartCmd) { + sessionStart = append(sessionStart, TraeHook{Name: "entire-session-start", Type: "command", Command: sessionStartCmd}) + count++ + } + if !hookCommandExists(sessionEnd, sessionEndCmd) { + sessionEnd = append(sessionEnd, TraeHook{Name: "entire-session-end", Type: "command", Command: sessionEndCmd}) + count++ + } + if !hookCommandExists(beforeAgent, beforeAgentCmd) { + beforeAgent = append(beforeAgent, TraeHook{Name: "entire-before-agent", Type: "command", Command: beforeAgentCmd}) + count++ + } + if !hookCommandExists(afterAgent, afterAgentCmd) { + afterAgent = append(afterAgent, TraeHook{Name: "entire-after-agent", Type: "command", Command: afterAgentCmd}) + count++ + } + if !hookCommandExists(beforeModel, beforeModelCmd) { + beforeModel = append(beforeModel, TraeHook{Name: "entire-before-model", Type: "command", Command: beforeModelCmd}) + count++ + } + if !hookCommandExists(afterModel, afterModelCmd) { + afterModel = append(afterModel, TraeHook{Name: "entire-after-model", Type: "command", Command: afterModelCmd}) + count++ + } + if !hookCommandExists(beforeToolSelection, beforeToolSelectionCmd) { + beforeToolSelection = append(beforeToolSelection, TraeHook{Name: "entire-before-tool-selection", Type: "command", Command: beforeToolSelectionCmd}) + count++ + } + if !hookCommandExists(preTool, preToolCmd) { + preTool = append(preTool, TraeHook{Name: "entire-pre-tool", Type: "command", Command: preToolCmd}) + count++ + } + if !hookCommandExists(afterTool, afterToolCmd) { + afterTool = append(afterTool, TraeHook{Name: "entire-after-tool", Type: "command", Command: afterToolCmd}) + count++ + } + if !hookCommandExists(preCompress, preCompressCmd) { + preCompress = append(preCompress, TraeHook{Name: "entire-pre-compress", Type: "command", Command: preCompressCmd}) + count++ + } + if !hookCommandExists(notification, notificationCmd) { + notification = append(notification, TraeHook{Name: "entire-notification", Type: "command", Command: notificationCmd}) + count++ + } + + if count == 0 { + return 0, nil // All hooks already installed + } + + // Marshal modified hook types back to rawHooks + marshalHookType(rawHooks, "SessionStart", sessionStart) + marshalHookType(rawHooks, "SessionEnd", sessionEnd) + marshalHookType(rawHooks, "BeforeAgent", beforeAgent) + marshalHookType(rawHooks, "AfterAgent", afterAgent) + marshalHookType(rawHooks, "BeforeModel", beforeModel) + marshalHookType(rawHooks, "AfterModel", afterModel) + marshalHookType(rawHooks, "BeforeToolSelection", beforeToolSelection) + marshalHookType(rawHooks, "PreTool", preTool) + marshalHookType(rawHooks, "AfterTool", afterTool) + marshalHookType(rawHooks, "PreCompress", preCompress) + marshalHookType(rawHooks, "Notification", notification) + + // Marshal hooks and update raw settings + hooksJSON, err := json.Marshal(rawHooks) + if err != nil { + return 0, fmt.Errorf("failed to marshal hooks: %w", err) + } + rawSettings["hooks"] = hooksJSON + + // Set hooksConfig.enabled = true (required for Trae Agent to execute hooks) + hooksConfig := TraeHooksConfig{Enabled: true} + hooksConfigJSON, err := json.Marshal(hooksConfig) + if err != nil { + return 0, fmt.Errorf("failed to marshal hooksConfig: %w", err) + } + rawSettings["hooksConfig"] = hooksConfigJSON + + // Write back to file + if err := os.MkdirAll(filepath.Dir(settingsPath), 0o750); err != nil { + return 0, fmt.Errorf("failed to create .trae directory: %w", err) + } + + output, err := jsonutil.MarshalIndentWithNewline(rawSettings, "", " ") + if err != nil { + return 0, fmt.Errorf("failed to marshal settings: %w", err) + } + + if err := os.WriteFile(settingsPath, output, 0o600); err != nil { + return 0, fmt.Errorf("failed to write settings.json: %w", err) + } + + return count, nil +} + +// parseHookType parses a specific hook type from rawHooks into the target slice. +// Silently ignores parse errors (leaves target unchanged). +func parseHookType(rawHooks map[string]json.RawMessage, hookType string, target interface{}) { + if data, ok := rawHooks[hookType]; ok { + //nolint:errcheck,gosec // Intentionally ignoring parse errors + json.Unmarshal(data, target) + } +} + +// marshalHookType marshals a hook type back to rawHooks. +// If the slice is empty, removes the key from rawHooks. +func marshalHookType(rawHooks map[string]json.RawMessage, hookType string, hooks interface{}) { + // Check if hooks is empty + var isEmpty bool + switch h := hooks.(type) { + case []TraeHook: + isEmpty = len(h) == 0 + default: + isEmpty = true + } + + if isEmpty { + delete(rawHooks, hookType) + return + } + + data, err := json.Marshal(hooks) + if err != nil { + return // Silently ignore marshal errors + } + rawHooks[hookType] = data +} + +// UninstallHooks removes Entire hooks from Trae Agent settings. +func (t *TraeAgent) UninstallHooks(ctx context.Context) error { + // Use repo root to find .trae directory when run from a subdirectory + repoRoot, err := paths.WorktreeRoot(ctx) + if err != nil { + repoRoot = "." // Fallback to CWD if not in a git repo + } + settingsPath := filepath.Join(repoRoot, ".trae", TraeSettingsFileName) + data, err := os.ReadFile(settingsPath) //nolint:gosec // path is constructed safely + if err != nil { + return nil //nolint:nilerr // No settings file means nothing to uninstall + } + + var rawSettings map[string]json.RawMessage + if err := json.Unmarshal(data, &rawSettings); err != nil { + return fmt.Errorf("failed to parse settings.json: %w", err) + } + + // rawHooks preserves unknown hook types + var rawHooks map[string]json.RawMessage + if hooksRaw, ok := rawSettings["hooks"]; ok { + if err := json.Unmarshal(hooksRaw, &rawHooks); err != nil { + return fmt.Errorf("failed to parse hooks: %w", err) + } + } + if rawHooks == nil { + rawHooks = make(map[string]json.RawMessage) + } + + // Parse all hook types + var sessionStart, sessionEnd, beforeAgent, afterAgent, beforeModel, afterModel []TraeHook + var beforeToolSelection, preTool, afterTool, preCompress, notification []TraeHook + + parseHookType(rawHooks, "SessionStart", &sessionStart) + parseHookType(rawHooks, "SessionEnd", &sessionEnd) + parseHookType(rawHooks, "BeforeAgent", &beforeAgent) + parseHookType(rawHooks, "AfterAgent", &afterAgent) + parseHookType(rawHooks, "BeforeModel", &beforeModel) + parseHookType(rawHooks, "AfterModel", &afterModel) + parseHookType(rawHooks, "BeforeToolSelection", &beforeToolSelection) + parseHookType(rawHooks, "PreTool", &preTool) + parseHookType(rawHooks, "AfterTool", &afterTool) + parseHookType(rawHooks, "PreCompress", &preCompress) + parseHookType(rawHooks, "Notification", ¬ification) + + // Remove Entire hooks from all hook types + sessionStart = removeEntireHooks(sessionStart) + sessionEnd = removeEntireHooks(sessionEnd) + beforeAgent = removeEntireHooks(beforeAgent) + afterAgent = removeEntireHooks(afterAgent) + beforeModel = removeEntireHooks(beforeModel) + afterModel = removeEntireHooks(afterModel) + beforeToolSelection = removeEntireHooks(beforeToolSelection) + preTool = removeEntireHooks(preTool) + afterTool = removeEntireHooks(afterTool) + preCompress = removeEntireHooks(preCompress) + notification = removeEntireHooks(notification) + + // Marshal modified hook types back to rawHooks + marshalHookType(rawHooks, "SessionStart", sessionStart) + marshalHookType(rawHooks, "SessionEnd", sessionEnd) + marshalHookType(rawHooks, "BeforeAgent", beforeAgent) + marshalHookType(rawHooks, "AfterAgent", afterAgent) + marshalHookType(rawHooks, "BeforeModel", beforeModel) + marshalHookType(rawHooks, "AfterModel", afterModel) + marshalHookType(rawHooks, "BeforeToolSelection", beforeToolSelection) + marshalHookType(rawHooks, "PreTool", preTool) + marshalHookType(rawHooks, "AfterTool", afterTool) + marshalHookType(rawHooks, "PreCompress", preCompress) + marshalHookType(rawHooks, "Notification", notification) + + // Marshal hooks back (preserving unknown hook types) + if len(rawHooks) > 0 { + hooksJSON, err := json.Marshal(rawHooks) + if err != nil { + return fmt.Errorf("failed to marshal hooks: %w", err) + } + rawSettings["hooks"] = hooksJSON + } else { + delete(rawSettings, "hooks") + } + + // Write back + output, err := jsonutil.MarshalIndentWithNewline(rawSettings, "", " ") + if err != nil { + return fmt.Errorf("failed to marshal settings: %w", err) + } + if err := os.WriteFile(settingsPath, output, 0o600); err != nil { + return fmt.Errorf("failed to write settings.json: %w", err) + } + return nil +} + +// AreHooksInstalled checks if Entire hooks are installed. +func (t *TraeAgent) AreHooksInstalled(ctx context.Context) bool { + // Use repo root to find .trae directory when run from a subdirectory + repoRoot, err := paths.WorktreeRoot(ctx) + if err != nil { + repoRoot = "." // Fallback to CWD if not in a git repo + } + settingsPath := filepath.Join(repoRoot, ".trae", TraeSettingsFileName) + data, err := os.ReadFile(settingsPath) //nolint:gosec // path is constructed safely + if err != nil { + return false + } + + var settings TraeSettings + if err := json.Unmarshal(data, &settings); err != nil { + return false + } + + // Check for at least one of our hooks + return hookCommandExists(settings.Hooks.SessionStart, "entire hooks trae-agent session-start") || + hookCommandExists(settings.Hooks.SessionStart, "go run ${TRAE_PROJECT_DIR}/cmd/entire/main.go hooks trae-agent session-start") +} + +// ParseHookEvent translates a Trae Agent hook into a normalized lifecycle Event. +// Returns nil if the hook has no lifecycle significance. +func (t *TraeAgent) ParseHookEvent(_ context.Context, hookName string, stdin io.Reader) (*agent.Event, error) { + switch hookName { + case HookNameSessionStart: + return t.parseSessionStart(stdin) + case HookNameBeforeAgent: + return t.parseBeforeAgent(stdin) + case HookNameAfterAgent: + return t.parseAfterAgent(stdin) + case HookNameSessionEnd: + return t.parseSessionEnd(stdin) + default: + return nil, nil //nolint:nilnil // Unknown hooks have no lifecycle action + } +} + +func (t *TraeAgent) parseSessionStart(stdin io.Reader) (*agent.Event, error) { + input, err := t.ParseHookInput(agent.HookSessionStart, stdin) + if err != nil { + return nil, err + } + return &agent.Event{ + Type: agent.SessionStart, + SessionID: input.SessionID, + SessionRef: input.SessionRef, + Timestamp: time.Now(), + }, nil +} + +func (t *TraeAgent) parseBeforeAgent(stdin io.Reader) (*agent.Event, error) { + input, err := t.ParseHookInput(agent.HookBeforeAgent, stdin) + if err != nil { + return nil, err + } + prompt, _ := input.RawData["prompt"].(string) + return &agent.Event{ + Type: agent.TurnStart, + SessionID: input.SessionID, + SessionRef: input.SessionRef, + Prompt: prompt, + Timestamp: time.Now(), + }, nil +} + +func (t *TraeAgent) parseAfterAgent(stdin io.Reader) (*agent.Event, error) { + input, err := t.ParseHookInput(agent.HookAfterAgent, stdin) + if err != nil { + return nil, err + } + return &agent.Event{ + Type: agent.TurnEnd, + SessionID: input.SessionID, + SessionRef: input.SessionRef, + Timestamp: time.Now(), + }, nil +} + +func (t *TraeAgent) parseSessionEnd(stdin io.Reader) (*agent.Event, error) { + input, err := t.ParseHookInput(agent.HookSessionEnd, stdin) + if err != nil { + return nil, err + } + return &agent.Event{ + Type: agent.SessionEnd, + SessionID: input.SessionID, + SessionRef: input.SessionRef, + Timestamp: time.Now(), + }, nil +} + +// GetSupportedHooks returns the hook types Trae Agent supports. +func (t *TraeAgent) GetSupportedHooks() []agent.HookType { + return []agent.HookType{ + agent.HookSessionStart, + agent.HookSessionEnd, + agent.HookBeforeAgent, + agent.HookAfterAgent, + agent.HookBeforeModel, + agent.HookAfterModel, + agent.HookBeforeToolSelection, + agent.HookPreTool, + agent.HookAfterTool, + agent.HookPreCompress, + agent.HookNotification, + } +} + +// Helper functions for hook management + +// TraeHook represents a single hook configuration in Trae Agent + +type TraeHook struct { + Name string `json:"name"` + Type string `json:"type"` + Command string `json:"command"` +} + +// TraeHooks represents all hook configurations in Trae Agent + +type TraeHooks struct { + SessionStart []TraeHook `json:"SessionStart,omitempty"` + SessionEnd []TraeHook `json:"SessionEnd,omitempty"` + BeforeAgent []TraeHook `json:"BeforeAgent,omitempty"` + AfterAgent []TraeHook `json:"AfterAgent,omitempty"` + BeforeModel []TraeHook `json:"BeforeModel,omitempty"` + AfterModel []TraeHook `json:"AfterModel,omitempty"` + BeforeToolSelection []TraeHook `json:"BeforeToolSelection,omitempty"` + PreTool []TraeHook `json:"PreTool,omitempty"` + AfterTool []TraeHook `json:"AfterTool,omitempty"` + PreCompress []TraeHook `json:"PreCompress,omitempty"` + Notification []TraeHook `json:"Notification,omitempty"` +} + +// TraeSettings represents the complete Trae Agent settings structure + +type TraeSettings struct { + HooksConfig TraeHooksConfig `json:"hooksConfig,omitempty"` + Hooks TraeHooks `json:"hooks,omitempty"` +} + +// TraeHooksConfig represents the hooks configuration settings +type TraeHooksConfig struct { + Enabled bool `json:"enabled,omitempty"` +} + +func hookCommandExists(hooks []TraeHook, command string) bool { + for _, hook := range hooks { + if hook.Command == command { + return true + } + } + return false +} + +// isEntireHook checks if a command is an Entire hook (old or new format) +func isEntireHook(command string) bool { + for _, prefix := range entireHookPrefixes { + if strings.HasPrefix(command, prefix) { + return true + } + } + return false +} + +// removeEntireHooks removes all Entire hooks from a list of hooks +func removeEntireHooks(hooks []TraeHook) []TraeHook { + result := make([]TraeHook, 0, len(hooks)) + for _, hook := range hooks { + if !isEntireHook(hook.Command) { + result = append(result, hook) + } + } + return result +} diff --git a/cmd/entire/cli/agent/traeagent/traeagent.go b/cmd/entire/cli/agent/traeagent/traeagent.go new file mode 100644 index 000000000..2de3fe9b4 --- /dev/null +++ b/cmd/entire/cli/agent/traeagent/traeagent.go @@ -0,0 +1,347 @@ +// Package traeagent implements the Agent interface for Trae Agent. +package traeagent + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "os" + "path/filepath" + "time" + + "github.com/entireio/cli/cmd/entire/cli/agent" + "github.com/entireio/cli/cmd/entire/cli/agent/types" + "github.com/entireio/cli/cmd/entire/cli/paths" +) + +//nolint:gochecknoinits // Agent self-registration is the intended pattern +func init() { + agent.Register(agent.AgentNameTraeAgent, NewTraeAgent) +} + +// TraeAgent implements the Agent interface for Trae Agent. +type TraeAgent struct{} + +// NewTraeAgent creates a new Trae Agent instance. +func NewTraeAgent() agent.Agent { + return &TraeAgent{} +} + +// Name returns the agent registry key. +func (t *TraeAgent) Name() types.AgentName { + return agent.AgentNameTraeAgent +} + +// Type returns the agent type identifier. +func (t *TraeAgent) Type() types.AgentType { + return agent.AgentTypeTraeAgent +} + +// Description returns a human-readable description. +func (t *TraeAgent) Description() string { + return "Trae Agent - ByteDance's LLM-based software engineering agent" +} + +// IsPreview returns whether the agent integration is in preview. +func (t *TraeAgent) IsPreview() bool { + return true +} + +// DetectPresence checks if Trae Agent is configured in the repository. +func (t *TraeAgent) DetectPresence(ctx context.Context) (bool, error) { + // Get repo root to check for .trae directory + repoRoot, err := paths.WorktreeRoot(ctx) + if err != nil { + // Not in a git repo, fall back to CWD-relative check + repoRoot = "." + } + + // Check for .trae directory + traeDir := filepath.Join(repoRoot, ".trae") + if _, err := os.Stat(traeDir); err == nil { + return true, nil + } + // Check for .trae/settings.json or trae_config.yaml + settingsFile := filepath.Join(repoRoot, ".trae", "settings.json") + if _, err := os.Stat(settingsFile); err == nil { + return true, nil + } + configFile := filepath.Join(repoRoot, "trae_config.yaml") + if _, err := os.Stat(configFile); err == nil { + return true, nil + } + return false, nil +} + +// GetHookConfigPath returns the path to Trae's hook config file. +func (t *TraeAgent) GetHookConfigPath() string { + return ".trae/settings.json" +} + +// SupportsHooks returns true as Trae Agent supports lifecycle hooks. +func (t *TraeAgent) SupportsHooks() bool { + return true +} + +// ParseHookInput parses Trae Agent hook input from stdin. +func (t *TraeAgent) ParseHookInput(hookType agent.HookType, reader io.Reader) (*agent.HookInput, error) { + data, err := io.ReadAll(reader) + if err != nil { + return nil, fmt.Errorf("failed to read input: %w", err) + } + + if len(data) == 0 { + return nil, errors.New("empty input") + } + + input := &agent.HookInput{ + HookType: hookType, + Timestamp: time.Now(), + RawData: make(map[string]interface{}), + } + + // Parse based on hook type + switch hookType { + case agent.HookSessionStart: + var raw sessionStartRaw + if err := json.Unmarshal(data, &raw); err != nil { + return nil, fmt.Errorf("failed to parse session start: %w", err) + } + input.SessionID = raw.SessionID + input.SessionRef = raw.TrajectoryPath + + case agent.HookSessionEnd: + var raw sessionEndRaw + if err := json.Unmarshal(data, &raw); err != nil { + return nil, fmt.Errorf("failed to parse session end: %w", err) + } + input.SessionID = raw.SessionID + input.SessionRef = raw.TrajectoryPath + + case agent.HookPreToolUse: + var raw preToolUseRaw + if err := json.Unmarshal(data, &raw); err != nil { + return nil, fmt.Errorf("failed to parse pre-tool input: %w", err) + } + input.SessionID = raw.SessionID + input.SessionRef = raw.TrajectoryPath + input.ToolName = raw.ToolName + input.ToolUseID = raw.ToolUseID + input.ToolInput = data + + case agent.HookPostToolUse: + var raw postToolUseRaw + if err := json.Unmarshal(data, &raw); err != nil { + return nil, fmt.Errorf("failed to parse post-tool input: %w", err) + } + input.SessionID = raw.SessionID + input.SessionRef = raw.TrajectoryPath + input.ToolName = raw.ToolName + input.ToolUseID = raw.ToolUseID + input.ToolInput = data + input.ToolResponse = data + + case agent.HookBeforeModel: + var raw preModelRaw + if err := json.Unmarshal(data, &raw); err != nil { + return nil, fmt.Errorf("failed to parse pre-model input: %w", err) + } + input.SessionID = raw.SessionID + input.SessionRef = raw.TrajectoryPath + input.RawData["model_name"] = raw.ModelName + input.RawData["prompt"] = raw.Prompt + + case agent.HookAfterModel: + var raw postModelRaw + if err := json.Unmarshal(data, &raw); err != nil { + return nil, fmt.Errorf("failed to parse post-model input: %w", err) + } + input.SessionID = raw.SessionID + input.SessionRef = raw.TrajectoryPath + input.RawData["model_name"] = raw.ModelName + input.RawData["response"] = raw.Response + input.RawData["token_usage"] = raw.TokenUsage + + case agent.HookPreCompress: + var raw preCompressRaw + if err := json.Unmarshal(data, &raw); err != nil { + return nil, fmt.Errorf("failed to parse pre-compress input: %w", err) + } + input.SessionID = raw.SessionID + input.SessionRef = raw.TrajectoryPath + input.RawData["context"] = raw.Context + + case agent.HookNotification: + var raw notificationRaw + if err := json.Unmarshal(data, &raw); err != nil { + return nil, fmt.Errorf("failed to parse notification input: %w", err) + } + input.SessionID = raw.SessionID + input.SessionRef = raw.TrajectoryPath + input.RawData["notification_type"] = raw.NotificationType + input.RawData["message"] = raw.Message + } + + return input, nil +} + +// GetSessionID extracts the session ID from hook input. +func (t *TraeAgent) GetSessionID(input *agent.HookInput) string { + return input.SessionID +} + +// ResolveSessionFile returns the path to a Trae Agent session file. +func (t *TraeAgent) ResolveSessionFile(sessionDir, agentSessionID string) string { + return filepath.Join(sessionDir, agentSessionID+"_trajectory.json") +} + +// ProtectedDirs returns directories that Trae Agent uses for config/state. +func (t *TraeAgent) ProtectedDirs() []string { return []string{".trae"} } + +// GetSessionDir returns the directory where Trae Agent stores session transcripts. +func (t *TraeAgent) GetSessionDir(repoPath string) (string, error) { + // Check for test environment override + if override := os.Getenv("ENTIRE_TEST_TRAE_PROJECT_DIR"); override != "" { + return override, nil + } + + homeDir, err := os.UserHomeDir() + if err != nil { + return "", fmt.Errorf("failed to get home directory: %w", err) + } + + return filepath.Join(homeDir, ".trae", "projects", filepath.Base(repoPath)), nil +} + +// ReadSession reads a session from Trae Agent's storage (JSON trajectory file). +func (t *TraeAgent) ReadSession(input *agent.HookInput) (*agent.AgentSession, error) { + if input.SessionRef == "" { + return nil, errors.New("session reference (trajectory path) is required") + } + + // Read the raw JSON trajectory file + data, err := os.ReadFile(input.SessionRef) + if err != nil { + return nil, fmt.Errorf("failed to read trajectory: %w", err) + } + + // Parse to extract computed fields + modifiedFiles, err := ExtractModifiedFiles(data) + if err != nil { + return nil, fmt.Errorf("failed to extract modified files: %w", err) + } + + return &agent.AgentSession{ + SessionID: input.SessionID, + AgentName: t.Name(), + SessionRef: input.SessionRef, + StartTime: time.Now(), + NativeData: data, + ModifiedFiles: modifiedFiles, + }, nil +} + +// WriteSession writes a session to Trae Agent's storage (JSON trajectory file). +func (t *TraeAgent) WriteSession(_ context.Context, session *agent.AgentSession) error { + if session == nil { + return errors.New("session is nil") + } + + // Verify this session belongs to Trae Agent + if session.AgentName != "" && session.AgentName != t.Name() { + return fmt.Errorf("session belongs to agent %q, not %q", session.AgentName, t.Name()) + } + + if session.SessionRef == "" { + return errors.New("session reference (trajectory path) is required") + } + + if len(session.NativeData) == 0 { + return errors.New("session has no native data to write") + } + + // Write the raw JSON data + if err := os.WriteFile(session.SessionRef, session.NativeData, 0o600); err != nil { + return fmt.Errorf("failed to write trajectory: %w", err) + } + + return nil +} + +// FormatResumeCommand returns the command to resume a Trae Agent session. +func (t *TraeAgent) FormatResumeCommand(sessionID string) string { + return "trae-cli interactive --resume " + sessionID +} + +// ReadTranscript reads the raw JSON trajectory bytes for a session. +func (t *TraeAgent) ReadTranscript(sessionRef string) ([]byte, error) { + data, err := os.ReadFile(sessionRef) //nolint:gosec // Path comes from agent hook input + if err != nil { + return nil, fmt.Errorf("failed to read trajectory: %w", err) + } + return data, nil +} + +// ExtractModifiedFiles extracts modified files from Trae Agent trajectory data. +func ExtractModifiedFiles(data []byte) ([]string, error) { + return extractModifiedFiles(data) +} + +// Raw data structures for parsing hooks + +type sessionStartRaw struct { + SessionID string `json:"session_id"` + TrajectoryPath string `json:"trajectory_path"` +} + +type sessionEndRaw struct { + SessionID string `json:"session_id"` + TrajectoryPath string `json:"trajectory_path"` +} + +type preToolUseRaw struct { + SessionID string `json:"session_id"` + TrajectoryPath string `json:"trajectory_path"` + ToolName string `json:"tool_name"` + ToolUseID string `json:"tool_use_id"` + ToolInput json.RawMessage `json:"tool_input"` +} + +type postToolUseRaw struct { + SessionID string `json:"session_id"` + TrajectoryPath string `json:"trajectory_path"` + ToolName string `json:"tool_name"` + ToolUseID string `json:"tool_use_id"` + ToolInput json.RawMessage `json:"tool_input"` + ToolResponse json.RawMessage `json:"tool_response"` +} + +type preModelRaw struct { + SessionID string `json:"session_id"` + TrajectoryPath string `json:"trajectory_path"` + ModelName string `json:"model_name"` + Prompt string `json:"prompt"` +} + +type postModelRaw struct { + SessionID string `json:"session_id"` + TrajectoryPath string `json:"trajectory_path"` + ModelName string `json:"model_name"` + Response string `json:"response"` + TokenUsage json.RawMessage `json:"token_usage"` +} + +type preCompressRaw struct { + SessionID string `json:"session_id"` + TrajectoryPath string `json:"trajectory_path"` + Context string `json:"context"` +} + +type notificationRaw struct { + SessionID string `json:"session_id"` + TrajectoryPath string `json:"trajectory_path"` + NotificationType string `json:"notification_type"` + Message string `json:"message"` +} diff --git a/cmd/entire/cli/agent/traeagent/traeagent_test.go b/cmd/entire/cli/agent/traeagent/traeagent_test.go new file mode 100644 index 000000000..4bbca6d05 --- /dev/null +++ b/cmd/entire/cli/agent/traeagent/traeagent_test.go @@ -0,0 +1,51 @@ +// Package traeagent implements the Agent interface for Trae Agent. +package traeagent + +import ( + "testing" + + "github.com/entireio/cli/cmd/entire/cli/agent" + "github.com/stretchr/testify/assert" +) + +func TestNewTraeAgent(t *testing.T) { + t.Parallel() + ag := NewTraeAgent() + assert.NotNil(t, ag) + assert.Equal(t, agent.AgentNameTraeAgent, ag.Name()) + assert.Equal(t, agent.AgentTypeTraeAgent, ag.Type()) + assert.Equal(t, "Trae Agent - ByteDance's LLM-based software engineering agent", ag.Description()) +} + +func TestTraeAgent_HookSupport(t *testing.T) { + t.Parallel() + ag := NewTraeAgent() + _, ok := ag.(agent.HookSupport) + assert.True(t, ok, "TraeAgent should implement HookSupport") +} + +func TestTraeAgent_ProtectedDirs(t *testing.T) { + t.Parallel() + ag := NewTraeAgent() + dirs := ag.ProtectedDirs() + assert.Equal(t, []string{".trae"}, dirs) +} + +func TestTraeAgent_HookNames(t *testing.T) { + t.Parallel() + ag := NewTraeAgent() + hookSupport, ok := ag.(agent.HookSupport) + assert.True(t, ok, "TraeAgent should implement HookSupport") + hookNames := hookSupport.HookNames() + assert.Contains(t, hookNames, HookNameSessionStart) + assert.Contains(t, hookNames, HookNameSessionEnd) + assert.Contains(t, hookNames, HookNameBeforeAgent) + assert.Contains(t, hookNames, HookNameAfterAgent) + assert.Contains(t, hookNames, HookNameBeforeModel) + assert.Contains(t, hookNames, HookNameAfterModel) + assert.Contains(t, hookNames, HookNameBeforeToolSelection) + assert.Contains(t, hookNames, HookNamePreTool) + assert.Contains(t, hookNames, HookNameAfterTool) + assert.Contains(t, hookNames, HookNamePreCompress) + assert.Contains(t, hookNames, HookNameNotification) +} diff --git a/cmd/entire/cli/agent/traeagent/transcript.go b/cmd/entire/cli/agent/traeagent/transcript.go new file mode 100644 index 000000000..88dc1da24 --- /dev/null +++ b/cmd/entire/cli/agent/traeagent/transcript.go @@ -0,0 +1,345 @@ +// Package traeagent implements the Agent interface for Trae Agent. +package traeagent + +import ( + "context" + "encoding/json" + "fmt" + "os" +) + +// Scanner buffer size for large transcript files (10MB) +const scannerBufferSize = 10 * 1024 * 1024 + +// TrajectoryEvent represents a single event in Trae Agent's trajectory + +type TrajectoryEvent struct { + Type string `json:"type"` + Timestamp string `json:"timestamp"` + EventID string `json:"event_id"` + Content json.RawMessage `json:"content"` + ToolName string `json:"tool_name,omitempty"` + ToolUseID string `json:"tool_use_id,omitempty"` + ToolInput json.RawMessage `json:"tool_input,omitempty"` + ToolOutput json.RawMessage `json:"tool_output,omitempty"` + ModelName string `json:"model_name,omitempty"` + Prompt string `json:"prompt,omitempty"` + Response string `json:"response,omitempty"` + TokenUsage json.RawMessage `json:"token_usage,omitempty"` +} + +// Trajectory represents the complete trajectory of a Trae Agent session +type Trajectory struct { + SessionID string `json:"session_id"` + StartTime string `json:"start_time"` + EndTime string `json:"end_time,omitempty"` + Events []TrajectoryEvent `json:"events"` +} + +// ParseTrajectory parses raw JSON content into a Trajectory object +func ParseTrajectory(data []byte) (*Trajectory, error) { + var trajectory Trajectory + if err := json.Unmarshal(data, &trajectory); err != nil { + return nil, fmt.Errorf("failed to parse trajectory: %w", err) + } + return &trajectory, nil +} + +// extractModifiedFiles extracts files modified by tool calls from trajectory +func extractModifiedFiles(data []byte) ([]string, error) { + trajectory, err := ParseTrajectory(data) + if err != nil { + return []string{}, nil // Return empty slice for now if parsing fails + } + + fileSet := make(map[string]bool) + var files []string + + for _, event := range trajectory.Events { + // Check for tool execution events + if event.Type == "tool_execution" || event.Type == "tool_result" { + // Check if it's a file modification tool + isModifyTool := false + for _, name := range FileModificationTools { + if event.ToolName == name { + isModifyTool = true + break + } + } + + if !isModifyTool { + continue + } + + // Try to extract file path from tool input + var toolInput struct { + FilePath string `json:"file_path,omitempty"` + Path string `json:"path,omitempty"` + } + + if err := json.Unmarshal(event.ToolInput, &toolInput); err == nil { + file := toolInput.FilePath + if file == "" { + file = toolInput.Path + } + + if file != "" && !fileSet[file] { + fileSet[file] = true + files = append(files, file) + } + } + } + } + + return files, nil +} + +// FileModificationTools lists the tools that modify files +var FileModificationTools = []string{ + "str_replace_based_edit_tool", + "edit_tool", + "write_file", + "update_file", + "delete_file", +} + +// TranscriptAnalyzer interface implementation + +// GetTranscriptPosition returns the current event count of a Trae Agent trajectory. +// Trae Agent uses JSON format with an events array, so position is the number of events. +func (t *TraeAgent) GetTranscriptPosition(path string) (int, error) { + if path == "" { + return 0, nil + } + + file, err := os.Open(path) //nolint:gosec // Path comes from Trae Agent trajectory location + if err != nil { + if os.IsNotExist(err) { + return 0, nil + } + return 0, fmt.Errorf("failed to open trajectory file: %w", err) + } + defer file.Close() + + // Read the entire file and parse it to get the event count + data, err := os.ReadFile(path) + if err != nil { + return 0, fmt.Errorf("failed to read trajectory file: %w", err) + } + + var trajectory Trajectory + if err := json.Unmarshal(data, &trajectory); err != nil { + return 0, fmt.Errorf("failed to parse trajectory: %w", err) + } + + return len(trajectory.Events), nil +} + +// ExtractModifiedFilesFromOffset extracts files modified since a given event index. +// For Trae Agent (JSON format), offset is the starting event index. +func (t *TraeAgent) ExtractModifiedFilesFromOffset(path string, startOffset int) (files []string, currentPosition int, err error) { + if path == "" { + return nil, 0, nil + } + + file, openErr := os.Open(path) //nolint:gosec // Path comes from Trae Agent trajectory location + if openErr != nil { + return nil, 0, fmt.Errorf("failed to open trajectory file: %w", openErr) + } + defer file.Close() + + // Read the entire file and parse it + data, err := os.ReadFile(path) + if err != nil { + return nil, 0, fmt.Errorf("failed to read trajectory file: %w", err) + } + + var trajectory Trajectory + if err := json.Unmarshal(data, &trajectory); err != nil { + return nil, 0, fmt.Errorf("failed to parse trajectory: %w", err) + } + + currentPosition = len(trajectory.Events) + if startOffset >= currentPosition { + return nil, currentPosition, nil + } + + // Extract events from startOffset onwards + relevantEvents := trajectory.Events[startOffset:] + + // Create a new trajectory with only relevant events + partialTrajectory := Trajectory{ + SessionID: trajectory.SessionID, + StartTime: trajectory.StartTime, + EndTime: trajectory.EndTime, + Events: relevantEvents, + } + + // Serialize partial trajectory and extract modified files + partialData, err := json.Marshal(partialTrajectory) + if err != nil { + return nil, currentPosition, fmt.Errorf("failed to marshal partial trajectory: %w", err) + } + + modifiedFiles, err := ExtractModifiedFiles(partialData) + if err != nil { + return nil, currentPosition, fmt.Errorf("failed to extract modified files: %w", err) + } + + return modifiedFiles, currentPosition, nil +} + +// TranscriptChunker interface implementation + +// ChunkTranscript splits a JSON trajectory into chunks if it exceeds maxSize. +// For JSON format, we split the events array into chunks while preserving valid JSON. +func (t *TraeAgent) ChunkTranscript(_ context.Context, content []byte, maxSize int) ([][]byte, error) { + // If content is smaller than maxSize, return as single chunk + if len(content) <= maxSize { + return [][]byte{content}, nil + } + + // Parse the trajectory to split events + var trajectory Trajectory + if err := json.Unmarshal(content, &trajectory); err != nil { + return nil, fmt.Errorf("failed to parse trajectory: %w", err) + } + + var chunks [][]byte + currentChunk := Trajectory{ + SessionID: trajectory.SessionID, + StartTime: trajectory.StartTime, + Events: []TrajectoryEvent{}, + } + + for _, event := range trajectory.Events { + // Add event to current chunk + currentChunk.Events = append(currentChunk.Events, event) + + // Check if current chunk exceeds maxSize + chunkData, err := json.Marshal(currentChunk) + if err != nil { + return nil, fmt.Errorf("failed to marshal chunk: %w", err) + } + + if len(chunkData) > maxSize { + // Remove the last event (it caused the overflow) + currentChunk.Events = currentChunk.Events[:len(currentChunk.Events)-1] + + // Marshal and add the chunk + finalChunkData, err := json.Marshal(currentChunk) + if err != nil { + return nil, fmt.Errorf("failed to marshal final chunk: %w", err) + } + chunks = append(chunks, finalChunkData) + + // Start a new chunk with the last event + currentChunk = Trajectory{ + SessionID: trajectory.SessionID, + StartTime: trajectory.StartTime, + Events: []TrajectoryEvent{event}, + } + } + } + + // Add the remaining chunk + if len(currentChunk.Events) > 0 { + chunkData, err := json.Marshal(currentChunk) + if err != nil { + return nil, fmt.Errorf("failed to marshal remaining chunk: %w", err) + } + chunks = append(chunks, chunkData) + } + + return chunks, nil +} + +// ReassembleTranscript combines chunks back into a single trajectory. +// For JSON format, we merge the events arrays from all chunks. +func (t *TraeAgent) ReassembleTranscript(chunks [][]byte) ([]byte, error) { + if len(chunks) == 0 { + return []byte("{}"), nil + } + + // Parse the first chunk to get the base trajectory + var mergedTrajectory Trajectory + if err := json.Unmarshal(chunks[0], &mergedTrajectory); err != nil { + return nil, fmt.Errorf("failed to parse first chunk: %w", err) + } + + // Merge events from remaining chunks + for i := 1; i < len(chunks); i++ { + var chunk Trajectory + if err := json.Unmarshal(chunks[i], &chunk); err != nil { + return nil, fmt.Errorf("failed to parse chunk %d: %w", i, err) + } + mergedTrajectory.Events = append(mergedTrajectory.Events, chunk.Events...) + } + + // Set end time from the last chunk if available + var lastChunk Trajectory + if err := json.Unmarshal(chunks[len(chunks)-1], &lastChunk); err == nil { + mergedTrajectory.EndTime = lastChunk.EndTime + } + + // Serialize the merged trajectory + mergedData, err := json.Marshal(mergedTrajectory) + if err != nil { + return nil, fmt.Errorf("failed to marshal merged trajectory: %w", err) + } + + return mergedData, nil +} + +// ExtractAllUserPrompts extracts all user prompts from a trajectory. +func ExtractAllUserPrompts(data []byte) ([]string, error) { + trajectory, err := ParseTrajectory(data) + if err != nil { + return nil, err + } + + var prompts []string + for _, event := range trajectory.Events { + if event.Type == "user_message" || event.Type == "user" { + var content struct { + Text string `json:"text"` + } + if err := json.Unmarshal(event.Content, &content); err == nil && content.Text != "" { + prompts = append(prompts, content.Text) + } + } + // Also check for prompt field directly in event + if event.Prompt != "" { + prompts = append(prompts, event.Prompt) + } + } + + return prompts, nil +} + +// ExtractLastAssistantMessage extracts the last assistant message from a trajectory. +func ExtractLastAssistantMessage(data []byte) (string, error) { + trajectory, err := ParseTrajectory(data) + if err != nil { + return "", err + } + + // Iterate in reverse to find the last assistant message + for i := len(trajectory.Events) - 1; i >= 0; i-- { + event := trajectory.Events[i] + if event.Type == "assistant_message" || event.Type == "assistant" { + var content struct { + Text string `json:"text"` + } + if err := json.Unmarshal(event.Content, &content); err == nil && content.Text != "" { + return content.Text, nil + } + } + // Also check for response field directly in event + if event.Response != "" { + return event.Response, nil + } + } + + return "", nil +} diff --git a/cmd/entire/cli/agent/types.go b/cmd/entire/cli/agent/types.go index 99f68b6d0..d50d8e3b1 100644 --- a/cmd/entire/cli/agent/types.go +++ b/cmd/entire/cli/agent/types.go @@ -6,12 +6,21 @@ import "time" type HookType string const ( - HookSessionStart HookType = "session_start" - HookSessionEnd HookType = "session_end" - HookUserPromptSubmit HookType = "user_prompt_submit" - HookStop HookType = "stop" - HookPreToolUse HookType = "pre_tool_use" - HookPostToolUse HookType = "post_tool_use" + HookSessionStart HookType = "session_start" + HookSessionEnd HookType = "session_end" + HookUserPromptSubmit HookType = "user_prompt_submit" + HookStop HookType = "stop" + HookPreToolUse HookType = "pre_tool_use" + HookPostToolUse HookType = "post_tool_use" + HookBeforeAgent HookType = "before_agent" + HookAfterAgent HookType = "after_agent" + HookBeforeModel HookType = "before_model" + HookAfterModel HookType = "after_model" + HookBeforeToolSelection HookType = "before_tool_selection" + HookPreTool HookType = "pre_tool" + HookAfterTool HookType = "after_tool" + HookPreCompress HookType = "pre_compress" + HookNotification HookType = "notification" ) // HookInput contains normalized data from hook callbacks diff --git a/cmd/entire/cli/hooks_cmd.go b/cmd/entire/cli/hooks_cmd.go index 74830ddbd..e5178e7e6 100644 --- a/cmd/entire/cli/hooks_cmd.go +++ b/cmd/entire/cli/hooks_cmd.go @@ -14,7 +14,9 @@ import ( _ "github.com/entireio/cli/cmd/entire/cli/agent/cursor" _ "github.com/entireio/cli/cmd/entire/cli/agent/factoryaidroid" _ "github.com/entireio/cli/cmd/entire/cli/agent/geminicli" + _ "github.com/entireio/cli/cmd/entire/cli/agent/iflow" _ "github.com/entireio/cli/cmd/entire/cli/agent/opencode" + _ "github.com/entireio/cli/cmd/entire/cli/agent/traeagent" _ "github.com/entireio/cli/cmd/entire/cli/agent/vogon" // support external agents diff --git a/cmd/entire/cli/setup.go b/cmd/entire/cli/setup.go index ecd0541c0..11cc6fa70 100644 --- a/cmd/entire/cli/setup.go +++ b/cmd/entire/cli/setup.go @@ -472,9 +472,24 @@ func runEnableInteractive(ctx context.Context, w io.Writer, agents []agent.Agent return fmt.Errorf("failed to clean up deselected agents: %w", err) } - // Setup agent hooks for all selected agents + // Load existing settings EARLY to get LocalDev for agent hook installation. + // This ensures re-running `entire enable` without --local-dev preserves existing localDev setting. + settings, err := LoadEntireSettings(ctx) + if err != nil { + // If we can't load, start with defaults + settings = &EntireSettings{} + } + // Merge opts flags into settings (opts takes precedence) + if opts.LocalDev { + settings.LocalDev = true + } + if opts.AbsoluteGitHookPath { + settings.AbsoluteGitHookPath = true + } + + // Setup agent hooks for all selected agents using merged settings.LocalDev for _, ag := range agents { - if _, err := setupAgentHooks(ctx, ag, opts.LocalDev, opts.ForceHooks); err != nil { + if _, err := setupAgentHooks(ctx, ag, settings.LocalDev, opts.ForceHooks); err != nil { return fmt.Errorf("failed to setup %s hooks: %w", ag.Type(), err) } } @@ -484,20 +499,8 @@ func runEnableInteractive(ctx context.Context, w io.Writer, agents []agent.Agent return fmt.Errorf("failed to setup .entire directory: %w", err) } - // Load existing settings to preserve other options (like strategy_options.push) - settings, err := LoadEntireSettings(ctx) - if err != nil { - // If we can't load, start with defaults - settings = &EntireSettings{} - } // Update the specific fields settings.Enabled = true - if opts.LocalDev { - settings.LocalDev = true - } - if opts.AbsoluteGitHookPath { - settings.AbsoluteGitHookPath = true - } opts.applyStrategyOptions(settings) @@ -958,24 +961,14 @@ func setupAgentHooksNonInteractive(ctx context.Context, w io.Writer, ag agent.Ag fmt.Fprintf(w, "Agent: %s\n\n", ag.Type()) - // Install agent hooks (agent hooks don't depend on settings) - installedHooks, err := hookAgent.InstallHooks(ctx, opts.LocalDev, opts.ForceHooks) - if err != nil { - return fmt.Errorf("failed to install hooks for %s: %w", agentName, err) - } - - // Setup .entire directory - if _, err := setupEntireDirectory(ctx); err != nil { - return fmt.Errorf("failed to setup .entire directory: %w", err) - } - - // Load existing settings to preserve other options (like strategy_options.push) + // Load existing settings EARLY to get LocalDev for agent hook installation. + // This ensures re-running `entire enable --agent X` without --local-dev preserves existing localDev setting. settings, err := LoadEntireSettings(ctx) if err != nil { // If we can't load, start with defaults settings = &EntireSettings{} } - settings.Enabled = true + // Merge opts flags into settings (opts takes precedence) if opts.LocalDev { settings.LocalDev = true } @@ -983,6 +976,20 @@ func setupAgentHooksNonInteractive(ctx context.Context, w io.Writer, ag agent.Ag settings.AbsoluteGitHookPath = true } + // Install agent hooks using merged settings.LocalDev + installedHooks, err := hookAgent.InstallHooks(ctx, settings.LocalDev, opts.ForceHooks) + if err != nil { + return fmt.Errorf("failed to install hooks for %s: %w", agentName, err) + } + + // Setup .entire directory + if _, err := setupEntireDirectory(ctx); err != nil { + return fmt.Errorf("failed to setup .entire directory: %w", err) + } + + // Update the specific fields + settings.Enabled = true + opts.applyStrategyOptions(settings) // Handle telemetry for non-interactive mode