diff --git a/cmd/main/agent.go b/cmd/main/agent.go index 00cc7c31..0d5e86fe 100644 --- a/cmd/main/agent.go +++ b/cmd/main/agent.go @@ -63,6 +63,7 @@ ADVANCED OPTIONS: --sandbox Sandbox config JSON for isolated execution --coding Coding config JSON for OpenCode integration --notify Enable notifications for this agent + --harness-config Agentic harness config JSON for multi-turn execution EXAMPLES: # Create a simple agent @@ -84,6 +85,13 @@ EXAMPLES: --prompt "Execute Python code in a sandbox." \ --sandbox '{"enabled":true,"image":"python:3.11"}' + # Create with agentic harness (Claude Agent SDK-like) + stn agent create coding-agent \ + --description "Multi-turn coding assistant" \ + --prompt "You are a coding assistant." \ + --tools read,write,bash,glob,grep,edit \ + --harness-config '{"max_steps":50,"sandbox":{"mode":"docker"}}' + # Create with structured output for CloudShip stn agent create cost-analyzer \ --description "Analyzes cloud costs" \ @@ -121,6 +129,7 @@ ADVANCED OPTIONS: --sandbox Update sandbox config JSON --coding Update coding config JSON --notify Update notification setting + --harness-config Update agentic harness config JSON EXAMPLES: # Update the prompt diff --git a/cmd/main/commands.go b/cmd/main/commands.go index 03a4e753..082b1d94 100644 --- a/cmd/main/commands.go +++ b/cmd/main/commands.go @@ -978,6 +978,70 @@ func runInit(cmd *cobra.Command, args []string) error { } } + // Set harness defaults (agentic execution harness) + if !viper.IsSet("harness.workspace.path") { + viper.Set("harness.workspace.path", "./workspace") + } + if !viper.IsSet("harness.workspace.mode") { + viper.Set("harness.workspace.mode", "host") + } + if !viper.IsSet("harness.compaction.enabled") { + viper.Set("harness.compaction.enabled", true) + } + if !viper.IsSet("harness.compaction.threshold") { + viper.Set("harness.compaction.threshold", "0.85") + } + if !viper.IsSet("harness.compaction.protect_tokens") { + viper.Set("harness.compaction.protect_tokens", 40000) + } + if !viper.IsSet("harness.git.auto_branch") { + viper.Set("harness.git.auto_branch", true) + } + if !viper.IsSet("harness.git.branch_prefix") { + viper.Set("harness.git.branch_prefix", "agent/") + } + if !viper.IsSet("harness.git.auto_commit") { + viper.Set("harness.git.auto_commit", false) + } + if !viper.IsSet("harness.git.require_approval") { + viper.Set("harness.git.require_approval", true) + } + if !viper.IsSet("harness.git.workflow_branch_strategy") { + viper.Set("harness.git.workflow_branch_strategy", "shared") + } + if !viper.IsSet("harness.nats.enabled") { + viper.Set("harness.nats.enabled", true) + } + if !viper.IsSet("harness.nats.kv_bucket") { + viper.Set("harness.nats.kv_bucket", "harness-state") + } + if !viper.IsSet("harness.nats.object_bucket") { + viper.Set("harness.nats.object_bucket", "harness-files") + } + if !viper.IsSet("harness.nats.max_file_size") { + viper.Set("harness.nats.max_file_size", "100MB") + } + if !viper.IsSet("harness.nats.ttl") { + viper.Set("harness.nats.ttl", "24h") + } + if !viper.IsSet("harness.permissions.external_directory") { + viper.Set("harness.permissions.external_directory", "deny") + } + + // Set sandbox defaults (Docker-based code execution) + if !viper.IsSet("sandbox.enabled") { + viper.Set("sandbox.enabled", false) + } + if !viper.IsSet("sandbox.code_mode_enabled") { + viper.Set("sandbox.code_mode_enabled", false) + } + if !viper.IsSet("sandbox.docker_image") { + viper.Set("sandbox.docker_image", "ubuntu:22.04") + } + if !viper.IsSet("sandbox.idle_timeout_minutes") { + viper.Set("sandbox.idle_timeout_minutes", 30) + } + // Generate encryption key if not already present if viper.GetString("encryption_key") == "" { encryptionKey, err := auth.GenerateAPIKey() @@ -997,6 +1061,14 @@ func runInit(cmd *cobra.Command, args []string) error { return fmt.Errorf("failed to write config file: %w", err) } + // Create harness workspace directory + harnessWorkspaceDir := filepath.Join(configDir, "workspace") + if err := os.MkdirAll(harnessWorkspaceDir, 0755); err != nil { + fmt.Printf(" Warning: Failed to create workspace directory: %v\n", err) + } else { + fmt.Printf(" Created harness workspace directory\n") + } + // Generate .gitignore for workspace if using custom config path if configPath != "" { gitignorePath := filepath.Join(configDir, ".gitignore") @@ -1006,6 +1078,9 @@ station.db station.db-shm station.db-wal +# Harness workspace (agent execution artifacts) +workspace/ + # Runtime variable resolution vars/ diff --git a/cmd/main/handlers/agent/handlers.go b/cmd/main/handlers/agent/handlers.go index 3c8d860e..16dba622 100644 --- a/cmd/main/handlers/agent/handlers.go +++ b/cmd/main/handlers/agent/handlers.go @@ -260,6 +260,7 @@ func (h *AgentHandler) RunAgentCreate(cmd *cobra.Command, args []string) error { sandboxConfig, _ := cmd.Flags().GetString("sandbox") codingConfig, _ := cmd.Flags().GetString("coding") notifyEnabled, _ := cmd.Flags().GetBool("notify") + harnessConfig, _ := cmd.Flags().GetString("harness-config") if sandboxConfig != "" { var sandboxMap map[string]interface{} @@ -273,6 +274,12 @@ func (h *AgentHandler) RunAgentCreate(cmd *cobra.Command, args []string) error { return fmt.Errorf("invalid --coding JSON: %v", err) } } + if harnessConfig != "" { + var harnessMap map[string]interface{} + if err := json.Unmarshal([]byte(harnessConfig), &harnessMap); err != nil { + return fmt.Errorf("invalid --harness-config JSON: %v", err) + } + } cfg, err := config.Load() if err != nil { @@ -341,8 +348,8 @@ func (h *AgentHandler) RunAgentCreate(cmd *cobra.Command, args []string) error { return fmt.Errorf("failed to create agent: %v", err) } - if sandboxConfig != "" || codingConfig != "" || notifyEnabled { - if err := exportService.ExportAgentWithConfigs(createdAgent.ID, app, appType, sandboxConfig, codingConfig, notifyEnabled); err != nil { + if sandboxConfig != "" || codingConfig != "" || harnessConfig != "" || notifyEnabled { + if err := exportService.ExportAgentWithAllConfigs(createdAgent.ID, app, appType, sandboxConfig, codingConfig, harnessConfig, notifyEnabled); err != nil { fmt.Printf("⚠️ Agent created but export failed: %v\n", err) } } @@ -397,6 +404,7 @@ func (h *AgentHandler) RunAgentUpdate(cmd *cobra.Command, args []string) error { sandboxConfig, _ := cmd.Flags().GetString("sandbox") codingConfig, _ := cmd.Flags().GetString("coding") notifyEnabled, _ := cmd.Flags().GetBool("notify") + harnessConfig, _ := cmd.Flags().GetString("harness-config") if sandboxConfig != "" { var sandboxMap map[string]interface{} @@ -410,6 +418,12 @@ func (h *AgentHandler) RunAgentUpdate(cmd *cobra.Command, args []string) error { return fmt.Errorf("invalid --coding JSON: %v", err) } } + if harnessConfig != "" { + var harnessMap map[string]interface{} + if err := json.Unmarshal([]byte(harnessConfig), &harnessMap); err != nil { + return fmt.Errorf("invalid --harness-config JSON: %v", err) + } + } cfg, err := config.Load() if err != nil { @@ -537,13 +551,17 @@ func (h *AgentHandler) RunAgentUpdate(cmd *cobra.Command, args []string) error { updates = append(updates, "notify") needsExport = true } + if cmd.Flags().Changed("harness-config") { + updates = append(updates, "harness-config") + needsExport = true + } if len(updates) == 0 { fmt.Printf("⚠️ No updates specified. Available flags:\n") fmt.Printf(" --prompt, --description, --max-steps, --tools, --output-schema\n") fmt.Printf(" --output-schema-preset, --app, --app-type\n") fmt.Printf(" --memory-topic, --memory-max-tokens\n") - fmt.Printf(" --sandbox, --coding, --notify\n") + fmt.Printf(" --sandbox, --coding, --notify, --harness-config\n") return nil } @@ -561,7 +579,7 @@ func (h *AgentHandler) RunAgentUpdate(cmd *cobra.Command, args []string) error { if needsExport { finalApp := agentConfig.App finalAppType := agentConfig.AppType - if err := exportService.ExportAgentWithConfigs(updatedAgent.ID, finalApp, finalAppType, sandboxConfig, codingConfig, notifyEnabled); err != nil { + if err := exportService.ExportAgentWithAllConfigs(updatedAgent.ID, finalApp, finalAppType, sandboxConfig, codingConfig, harnessConfig, notifyEnabled); err != nil { fmt.Printf("⚠️ Agent updated but export failed: %v\n", err) } } diff --git a/cmd/main/harness_repl.go b/cmd/main/harness_repl.go new file mode 100644 index 00000000..dabac46e --- /dev/null +++ b/cmd/main/harness_repl.go @@ -0,0 +1,650 @@ +package main + +import ( + "bufio" + "context" + "encoding/json" + "fmt" + "os" + "os/signal" + "path/filepath" + "strings" + "syscall" + "time" + + "github.com/spf13/cobra" + + "station/internal/config" + "station/internal/db" + "station/internal/db/repositories" + "station/internal/services" + "station/pkg/harness/memory" + "station/pkg/harness/session" + "station/pkg/harness/skills" + "station/pkg/models" +) + +var harnessCmd = &cobra.Command{ + Use: "harness", + Short: "Agentic harness commands", + Long: "Commands for working with the agentic harness system", +} + +var harnessReplCmd = &cobra.Command{ + Use: "repl", + Short: "Start interactive REPL for agent testing", + Long: `Start an interactive REPL (Read-Eval-Print-Loop) for testing agentic harness agents. + +The REPL provides a conversational interface where you can: +• Send tasks to the agent and see responses +• Track session state, steps, and token usage +• List available skills and tools +• Save and load sessions for later + +EXAMPLES: + # Start REPL for an agent + stn harness repl --agent myagent --env default + + # Resume a previous session + stn harness repl --session abc123 + + # Start with specific sandbox mode + stn harness repl --agent myagent --sandbox docker + +REPL COMMANDS: + /help Show all available commands + /status Show session status (steps, tokens, progress) + /history Show conversation history summary + /skills List available skills + /tools List available tools + /files List workspace files + /reset Reset session (start fresh) + /save Save session to file + /load Load session from file + /exit Exit REPL`, + RunE: runHarnessRepl, +} + +func init() { + harnessReplCmd.Flags().StringP("agent", "a", "", "Agent name to run") + harnessReplCmd.Flags().StringP("env", "e", "default", "Environment name") + harnessReplCmd.Flags().StringP("session", "s", "", "Session ID to resume") + harnessReplCmd.Flags().String("sandbox", "", "Sandbox mode (host, docker, modal, etc.)") + harnessReplCmd.Flags().StringSlice("skills", nil, "Additional skill source paths") + harnessReplCmd.Flags().Int("max-steps", 50, "Maximum steps per execution") + harnessReplCmd.Flags().String("timeout", "30m", "Execution timeout") + + harnessCmd.AddCommand(harnessReplCmd) +} + +// ReplSession holds the state for a REPL session +type ReplSession struct { + ID string + AgentID string + AgentName string + Task string + StepCount int + MessageCount int + TotalTokens int + CreatedAt time.Time + UpdatedAt time.Time + + // Session persistence + sessionManager *session.Manager + historyStore *session.HistoryStore + persistedPath string +} + +func runHarnessRepl(cmd *cobra.Command, args []string) error { + ctx, cancel := context.WithCancel(cmd.Context()) + defer cancel() + + // Handle Ctrl+C gracefully + sigChan := make(chan os.Signal, 1) + signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM) + go func() { + <-sigChan + fmt.Println("\n\nGoodbye!") + cancel() + os.Exit(0) + }() + + agentName, _ := cmd.Flags().GetString("agent") + envName, _ := cmd.Flags().GetString("env") + sessionID, _ := cmd.Flags().GetString("session") + sandboxMode, _ := cmd.Flags().GetString("sandbox") + skillSources, _ := cmd.Flags().GetStringSlice("skills") + maxSteps, _ := cmd.Flags().GetInt("max-steps") + timeoutStr, _ := cmd.Flags().GetString("timeout") + + timeout, err := time.ParseDuration(timeoutStr) + if err != nil { + timeout = 30 * time.Minute + } + + // Validate inputs + if agentName == "" && sessionID == "" { + return fmt.Errorf("either --agent or --session is required") + } + + // Load config + cfg, err := config.Load() + if err != nil { + return fmt.Errorf("failed to load config: %w", err) + } + + // Initialize database + database, err := db.New(cfg.DatabaseURL) + if err != nil { + return fmt.Errorf("failed to connect to database: %w", err) + } + defer func() { _ = database.Close() }() + + repos := repositories.New(database) + + // Get environment + env, err := repos.Environments.GetByName(envName) + if err != nil { + return fmt.Errorf("environment '%s' not found: %w", envName, err) + } + + // Get agent + var agent *models.Agent + if agentName != "" { + agents, err := repos.Agents.ListByEnvironment(env.ID) + if err != nil { + return fmt.Errorf("failed to list agents: %w", err) + } + + for _, a := range agents { + if a.Name == agentName { + agent = a + break + } + } + + if agent == nil { + return fmt.Errorf("agent '%s' not found in environment '%s'", agentName, envName) + } + } + + // Initialize skills middleware + var skillsMiddleware *skills.SkillsMiddleware + envPath := config.GetEnvironmentDir(envName) + skillSourcePaths := skills.DefaultSkillSources(envPath) + if len(skillSources) > 0 { + skillSourcePaths = append(skillSourcePaths, skillSources...) + } + skillsBackend := skills.NewFSBackend() + skillsMiddleware = skills.NewSkillsMiddleware(skillsBackend, skillSourcePaths) + + // Initialize memory middleware + memoryBackend := memory.NewFSBackend() + memorySources := memory.DefaultMemorySources(envPath) + memoryMiddleware := memory.NewMemoryMiddleware(memoryBackend, memorySources) + + // Initialize session persistence + workspacePath := cfg.Harness.Workspace.Path + if workspacePath == "" { + workspacePath = filepath.Join(cfg.Workspace, "workspace") + } + if !filepath.IsAbs(workspacePath) { + workspacePath = filepath.Join(cfg.Workspace, workspacePath) + } + sessionMgr := session.NewManager(workspacePath) + historyStore := session.NewHistoryStore(workspacePath) + + // Create or resume session + replSession := &ReplSession{ + ID: fmt.Sprintf("repl-%d", time.Now().Unix()), + AgentName: agentName, + CreatedAt: time.Now(), + UpdatedAt: time.Now(), + sessionManager: sessionMgr, + historyStore: historyStore, + } + if agent != nil { + replSession.AgentID = fmt.Sprintf("%d", agent.ID) + } + + // Resume or create persisted session + if sessionID != "" { + replSession.ID = sessionID + // Load existing session metadata + existingSession, err := sessionMgr.Get(ctx, sessionID) + if err == nil { + replSession.persistedPath = existingSession.Path + // Load history to get message count and tokens + history, err := historyStore.Load(sessionID) + if err == nil { + replSession.MessageCount = len(history.Messages) + replSession.TotalTokens = history.TotalTokens + } + fmt.Printf("Resuming session %s (%d messages)\n", sessionID, replSession.MessageCount) + } else { + fmt.Printf("Session %s not found, creating new...\n", sessionID) + } + } + + // Ensure session directory exists + if _, err := sessionMgr.GetOrCreate(ctx, replSession.ID, "", ""); err != nil { + fmt.Printf("Warning: failed to create session directory: %v\n", err) + } + + // Print banner + printReplBanner(replSession, sandboxMode, envName) + + // Main REPL loop + reader := bufio.NewReader(os.Stdin) + + for { + select { + case <-ctx.Done(): + return nil + default: + } + + fmt.Print(">>> ") + input, err := reader.ReadString('\n') + if err != nil { + break + } + + input = strings.TrimSpace(input) + if input == "" { + continue + } + + // Handle REPL commands + if strings.HasPrefix(input, "/") { + if handleReplCommand(ctx, replSession, skillsMiddleware, memoryMiddleware, input) { + return nil // /exit + } + continue + } + + // Execute agent task + fmt.Println() + result, err := executeReplTask(ctx, cfg, repos, env, agent, input, maxSteps, timeout, sandboxMode, skillsMiddleware, memoryMiddleware) + if err != nil { + fmt.Printf("Error: %v\n\n", err) + continue + } + + // Update session stats + replSession.StepCount += result.TotalSteps + replSession.TotalTokens += result.TotalTokens + replSession.MessageCount++ + replSession.UpdatedAt = time.Now() + + // Persist message history + if replSession.historyStore != nil { + newMessages := []session.StoredMessage{ + { + Role: "user", + Content: input, + Timestamp: time.Now(), + }, + { + Role: "assistant", + Content: result.Response, + Timestamp: time.Now(), + }, + } + if err := replSession.historyStore.Append(replSession.ID, newMessages); err != nil { + fmt.Printf("Warning: failed to save history: %v\n", err) + } + } + + // Print response + fmt.Println(result.Response) + fmt.Printf("\n[Steps: %d | Tokens: %d | %s]\n\n", + result.TotalSteps, result.TotalTokens, result.FinishReason) + } + + return nil +} + +func printReplBanner(session *ReplSession, sandboxMode, envName string) { + if sandboxMode == "" { + sandboxMode = "host" + } + + sessionDisplay := session.ID + if len(sessionDisplay) > 8 { + sessionDisplay = sessionDisplay[:8] + "..." + } + + fmt.Println() + fmt.Println("+=============================================+") + fmt.Println("| Station Agent REPL |") + fmt.Println("+=============================================+") + fmt.Printf("| Agent: %-33s |\n", session.AgentName) + fmt.Printf("| Session: %-33s |\n", sessionDisplay) + fmt.Printf("| Env: %-33s |\n", envName) + fmt.Printf("| Sandbox: %-33s |\n", sandboxMode) + fmt.Println("+---------------------------------------------+") + fmt.Println("| Commands: |") + fmt.Println("| /help - Show all commands |") + fmt.Println("| /status - Session status |") + fmt.Println("| /exit - Exit REPL |") + fmt.Println("+=============================================+") + fmt.Println() +} + +func handleReplCommand(ctx context.Context, session *ReplSession, skillsMW *skills.SkillsMiddleware, memoryMW *memory.MemoryMiddleware, input string) bool { + parts := strings.Fields(input) + cmdName := parts[0] + + switch cmdName { + case "/exit", "/quit", "/q": + fmt.Println("Goodbye!") + return true + + case "/help", "/h", "/?": + printReplHelp() + + case "/status": + printReplSessionStatus(session) + + case "/skills": + if skillsMW != nil { + loadedSkills, err := skillsMW.LoadSkills() + if err != nil { + fmt.Printf("Error loading skills: %v\n", err) + } else if len(loadedSkills) == 0 { + fmt.Println("No skills found") + } else { + fmt.Println("\nAvailable Skills:") + for _, skill := range loadedSkills { + fmt.Printf(" - %s: %s\n", skill.Name, skill.Description) + fmt.Printf(" Path: %s\n", skill.Path) + } + fmt.Println() + } + } else { + fmt.Println("Skills middleware not initialized") + } + + case "/memory": + if memoryMW != nil { + contents, err := memoryMW.LoadMemory() + if err != nil { + fmt.Printf("Error loading memory: %v\n", err) + } else if len(contents) == 0 { + fmt.Println("No memory sources found") + } else { + fmt.Println("\nLoaded Memory Sources:") + for source := range contents { + fmt.Printf(" - %s\n", source) + } + fmt.Println() + } + } else { + fmt.Println("Memory middleware not initialized") + } + + case "/tools": + fmt.Println("\nAvailable Tools:") + fmt.Println(" - bash: Execute shell commands") + fmt.Println(" - read_file: Read file contents") + fmt.Println(" - write_file: Write file contents") + fmt.Println(" - edit_file: Edit file contents") + fmt.Println(" - glob: Find files by pattern") + fmt.Println(" - grep: Search file contents") + fmt.Println(" (Additional MCP tools depend on environment configuration)") + fmt.Println() + + case "/reset": + // Clear persisted history before resetting + if session.historyStore != nil { + _ = session.historyStore.Clear(session.ID) + } + session.StepCount = 0 + session.TotalTokens = 0 + session.MessageCount = 0 + session.ID = fmt.Sprintf("repl-%d", time.Now().Unix()) + session.CreatedAt = time.Now() + session.UpdatedAt = time.Now() + // Create new session directory + if session.sessionManager != nil { + _, _ = session.sessionManager.GetOrCreate(ctx, session.ID, "", "") + } + fmt.Printf("Session reset. New session: %s\n\n", session.ID[:12]+"...") + + case "/history": + fmt.Printf("\nSession History:\n") + fmt.Printf(" Session ID: %s\n", session.ID) + fmt.Printf(" Total messages: %d\n", session.MessageCount) + fmt.Printf(" Total steps: %d\n", session.StepCount) + fmt.Printf(" Total tokens: %d\n", session.TotalTokens) + fmt.Printf(" Duration: %s\n", time.Since(session.CreatedAt).Round(time.Second)) + + // Show recent messages from persisted history + if session.historyStore != nil { + history, err := session.historyStore.Load(session.ID) + if err == nil && len(history.Messages) > 0 { + fmt.Printf("\n Recent messages (last 5):\n") + start := 0 + if len(history.Messages) > 10 { + start = len(history.Messages) - 10 + } + for i := start; i < len(history.Messages); i++ { + msg := history.Messages[i] + content := msg.Content + if len(content) > 60 { + content = content[:60] + "..." + } + fmt.Printf(" [%s] %s\n", msg.Role, content) + } + } + } + fmt.Println() + + case "/clear": + // Clear history but keep session + if session.historyStore != nil { + if err := session.historyStore.Clear(session.ID); err != nil { + fmt.Printf("Failed to clear history: %v\n", err) + } else { + session.MessageCount = 0 + session.TotalTokens = 0 + fmt.Println("History cleared.\n") + } + } else { + fmt.Println("History store not initialized") + } + + case "/save": + if len(parts) > 1 { + filename := parts[1] + err := saveReplSession(session, filename) + if err != nil { + fmt.Printf("Save failed: %v\n", err) + } else { + fmt.Printf("Session saved to: %s\n", filename) + } + } else { + fmt.Println("Usage: /save ") + } + + case "/load": + if len(parts) > 1 { + filename := parts[1] + loaded, err := loadReplSession(filename) + if err != nil { + fmt.Printf("Load failed: %v\n", err) + } else { + *session = *loaded + fmt.Printf("Session loaded: %s\n", session.ID[:12]+"...") + } + } else { + fmt.Println("Usage: /load ") + } + + case "/files": + fmt.Println("\nWorkspace file listing not implemented yet") + fmt.Println("Use: /tools to see available file operations\n") + + default: + fmt.Printf("Unknown command: %s (try /help)\n\n", cmdName) + } + + return false +} + +func printReplHelp() { + fmt.Println(` +REPL Commands: + /help, /h Show this help + /status Show session status (steps, tokens) + /history Show conversation history with recent messages + /clear Clear history but keep session + /skills List available skills + /memory List loaded memory sources + /tools List available tools + /files List workspace files + /reset Reset session (start fresh) + /save FILE Save session state to file + /load FILE Load session state from file + /exit, /q Exit REPL + +Any other input is sent to the agent as a task. +`) +} + +func printReplSessionStatus(session *ReplSession) { + fmt.Println() + fmt.Println("Session Status:") + fmt.Printf(" ID: %s\n", session.ID) + fmt.Printf(" Agent: %s\n", session.AgentName) + fmt.Printf(" Steps: %d\n", session.StepCount) + fmt.Printf(" Tokens: %d\n", session.TotalTokens) + fmt.Printf(" Messages: %d\n", session.MessageCount) + fmt.Printf(" Created: %s\n", session.CreatedAt.Format("2006-01-02 15:04:05")) + fmt.Printf(" Updated: %s\n", session.UpdatedAt.Format("2006-01-02 15:04:05")) + if session.persistedPath != "" { + fmt.Printf(" Path: %s\n", session.persistedPath) + } + fmt.Printf(" Persisted: %v\n", session.historyStore != nil) + fmt.Println() +} + +// ReplSessionExport is a JSON-serializable version of ReplSession +type ReplSessionExport struct { + ID string `json:"id"` + AgentID string `json:"agent_id"` + AgentName string `json:"agent_name"` + StepCount int `json:"step_count"` + MessageCount int `json:"message_count"` + TotalTokens int `json:"total_tokens"` + CreatedAt time.Time `json:"created_at"` + UpdatedAt time.Time `json:"updated_at"` +} + +func saveReplSession(sess *ReplSession, filename string) error { + export := ReplSessionExport{ + ID: sess.ID, + AgentID: sess.AgentID, + AgentName: sess.AgentName, + StepCount: sess.StepCount, + MessageCount: sess.MessageCount, + TotalTokens: sess.TotalTokens, + CreatedAt: sess.CreatedAt, + UpdatedAt: sess.UpdatedAt, + } + + data, err := json.MarshalIndent(export, "", " ") + if err != nil { + return fmt.Errorf("failed to marshal session: %w", err) + } + + if err := os.WriteFile(filename, data, 0644); err != nil { + return fmt.Errorf("failed to write file: %w", err) + } + + return nil +} + +func loadReplSession(filename string) (*ReplSession, error) { + data, err := os.ReadFile(filename) + if err != nil { + return nil, err + } + + var export ReplSessionExport + if err := json.Unmarshal(data, &export); err != nil { + return nil, fmt.Errorf("failed to parse session file: %w", err) + } + + sess := &ReplSession{ + ID: export.ID, + AgentID: export.AgentID, + AgentName: export.AgentName, + StepCount: export.StepCount, + MessageCount: export.MessageCount, + TotalTokens: export.TotalTokens, + CreatedAt: export.CreatedAt, + UpdatedAt: export.UpdatedAt, + } + + if sess.ID == "" { + sess.ID = fmt.Sprintf("repl-%d", time.Now().Unix()) + } + + return sess, nil +} + +// ExecutionResult represents the result of an agent execution +type ExecutionResult struct { + Response string + TotalSteps int + TotalTokens int + FinishReason string + Error error +} + +func executeReplTask( + ctx context.Context, + cfg *config.Config, + repos *repositories.Repositories, + env *models.Environment, + agent *models.Agent, + task string, + maxSteps int, + timeout time.Duration, + sandboxMode string, + skillsMW *skills.SkillsMiddleware, + memoryMW *memory.MemoryMiddleware, +) (*ExecutionResult, error) { + // Create execution context with timeout + execCtx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + + // Use the agent execution service + agentService := services.NewAgentService(repos) + + // Execute through the standard service which handles harness internally + // ExecuteAgent returns Message directly, no need to poll + message, err := agentService.ExecuteAgent(execCtx, agent.ID, task, nil) + if err != nil { + return nil, fmt.Errorf("execution failed: %w", err) + } + + // Build result from message + // Note: The simple Message struct doesn't include token/step counts + // These would need to be retrieved from the run record if needed + result := &ExecutionResult{ + Response: message.Content, + TotalSteps: 1, // Placeholder - could be retrieved from run record + TotalTokens: 0, // Placeholder - could be retrieved from run record + FinishReason: "completed", + } + + return result, nil +} + +// GetHarnessCmd returns the harness command for registration +func GetHarnessCmd() *cobra.Command { + return harnessCmd +} diff --git a/cmd/main/main.go b/cmd/main/main.go index 3b6b2bb9..740133e6 100644 --- a/cmd/main/main.go +++ b/cmd/main/main.go @@ -76,6 +76,8 @@ func init() { rootCmd.AddCommand(latticeCmd) rootCmd.AddCommand(githubCmd) rootCmd.AddCommand(secretsCmd) + rootCmd.AddCommand(sessionCmd) + rootCmd.AddCommand(harnessCmd) filesCmd.AddCommand(filesUploadCmd) filesCmd.AddCommand(filesDownloadCmd) @@ -279,6 +281,7 @@ func init() { agentCreateCmd.Flags().String("sandbox", "", "Sandbox config JSON (e.g., '{\"enabled\":true,\"image\":\"python:3.11\"}')") agentCreateCmd.Flags().String("coding", "", "Coding config JSON (e.g., '{\"enabled\":true,\"backend\":\"opencode\"}')") agentCreateCmd.Flags().Bool("notify", false, "Enable notifications for this agent") + agentCreateCmd.Flags().String("harness-config", "", "Agentic harness config JSON (e.g., '{\"max_steps\":50,\"sandbox\":{\"mode\":\"docker\"}}')") agentUpdateCmd.Flags().StringP("environment", "e", "default", "Environment the agent is in") agentUpdateCmd.Flags().StringP("prompt", "p", "", "New system prompt") @@ -294,6 +297,7 @@ func init() { agentUpdateCmd.Flags().String("sandbox", "", "Sandbox config JSON") agentUpdateCmd.Flags().String("coding", "", "Coding config JSON") agentUpdateCmd.Flags().Bool("notify", false, "Enable notifications") + agentUpdateCmd.Flags().String("harness-config", "", "Agentic harness config JSON") // Runs command flags runsListCmd.Flags().String("endpoint", "", "Station API endpoint (default: use local mode)") diff --git a/cmd/main/session_commands.go b/cmd/main/session_commands.go new file mode 100644 index 00000000..dc442004 --- /dev/null +++ b/cmd/main/session_commands.go @@ -0,0 +1,318 @@ +package main + +import ( + "context" + "fmt" + "os" + "path/filepath" + "strings" + "text/tabwriter" + "time" + + "station/internal/config" + "station/pkg/harness/session" + + "github.com/spf13/cobra" +) + +var ( + sessionCmd = &cobra.Command{ + Use: "session", + Short: "Manage harness workspace sessions", + Long: `Manage persistent workspace sessions for agentic harness executions.`, + } + + sessionListCmd = &cobra.Command{ + Use: "list", + Short: "List all sessions", + RunE: runSessionList, + } + + sessionInfoCmd = &cobra.Command{ + Use: "info ", + Short: "Show session details", + Args: cobra.ExactArgs(1), + RunE: runSessionInfo, + } + + sessionDeleteCmd = &cobra.Command{ + Use: "delete ", + Short: "Delete a session", + Args: cobra.ExactArgs(1), + RunE: runSessionDelete, + } + + sessionCleanupCmd = &cobra.Command{ + Use: "cleanup", + Short: "Remove stale sessions", + RunE: runSessionCleanup, + } + + sessionUnlockCmd = &cobra.Command{ + Use: "unlock ", + Short: "Force unlock a session", + Args: cobra.ExactArgs(1), + RunE: runSessionUnlock, + } +) + +var ( + cleanupOlderThan string + cleanupDryRun bool + deleteForce bool +) + +func init() { + sessionCleanupCmd.Flags().StringVar(&cleanupOlderThan, "older-than", "7d", "Delete sessions older than duration (e.g., 7d, 24h)") + sessionCleanupCmd.Flags().BoolVar(&cleanupDryRun, "dry-run", false, "Show what would be deleted without actually deleting") + + sessionDeleteCmd.Flags().BoolVar(&deleteForce, "force", false, "Force delete even if session is locked") + + sessionCmd.AddCommand(sessionListCmd) + sessionCmd.AddCommand(sessionInfoCmd) + sessionCmd.AddCommand(sessionDeleteCmd) + sessionCmd.AddCommand(sessionCleanupCmd) + sessionCmd.AddCommand(sessionUnlockCmd) +} + +func getSessionManager() *session.Manager { + cfg := config.GetLoadedConfig() + basePath := "./workspace" + if cfg != nil && cfg.Harness.Workspace.Path != "" { + basePath = cfg.Harness.Workspace.Path + } + if !filepath.IsAbs(basePath) { + if cfg != nil && cfg.Workspace != "" { + basePath = filepath.Join(cfg.Workspace, basePath) + } + } + return session.NewManager(basePath) +} + +func runSessionList(cmd *cobra.Command, args []string) error { + mgr := getSessionManager() + ctx := context.Background() + + sessions, err := mgr.List(ctx) + if err != nil { + return fmt.Errorf("failed to list sessions: %w", err) + } + + if len(sessions) == 0 { + fmt.Println("No sessions found.") + return nil + } + + w := tabwriter.NewWriter(os.Stdout, 0, 0, 2, ' ', 0) + fmt.Fprintln(w, "SESSION ID\tSTATUS\tLAST USED\tRUNS\tREPO\tSIZE") + fmt.Fprintln(w, "----------\t------\t---------\t----\t----\t----") + + for _, s := range sessions { + status := "ready" + if s.IsLocked { + status = fmt.Sprintf("locked (%s)", s.LockedBy) + } + + lastUsed := formatSessionDuration(time.Since(s.LastUsedAt)) + repo := s.RepoURL + if len(repo) > 40 { + repo = "..." + repo[len(repo)-37:] + } + if repo == "" { + repo = "-" + } + + size, _ := mgr.DiskUsage(s.ID) + sizeStr := formatSessionSize(size) + + fmt.Fprintf(w, "%s\t%s\t%s ago\t%d\t%s\t%s\n", + s.ID, status, lastUsed, s.TotalRuns, repo, sizeStr) + } + w.Flush() + + return nil +} + +func runSessionInfo(cmd *cobra.Command, args []string) error { + sessionID := args[0] + mgr := getSessionManager() + ctx := context.Background() + + s, err := mgr.Get(ctx, sessionID) + if err != nil { + return fmt.Errorf("session not found: %s", sessionID) + } + + size, _ := mgr.DiskUsage(s.ID) + + fmt.Printf("Session: %s\n", s.ID) + fmt.Printf("Path: %s\n", s.Path) + fmt.Printf("Created: %s\n", s.CreatedAt.Format(time.RFC3339)) + fmt.Printf("Last Used: %s (%s ago)\n", s.LastUsedAt.Format(time.RFC3339), formatSessionDuration(time.Since(s.LastUsedAt))) + fmt.Printf("Total Runs: %d\n", s.TotalRuns) + fmt.Printf("Disk Usage: %s\n", formatSessionSize(size)) + + if s.RepoURL != "" { + fmt.Printf("Repository: %s\n", s.RepoURL) + } + if s.Branch != "" { + fmt.Printf("Branch: %s\n", s.Branch) + } + + if s.IsLocked { + fmt.Printf("Status: LOCKED\n") + fmt.Printf("Locked By: %s\n", s.LockedBy) + fmt.Printf("Locked At: %s\n", s.LockedAt.Format(time.RFC3339)) + } else { + fmt.Printf("Status: Ready\n") + } + + return nil +} + +func runSessionDelete(cmd *cobra.Command, args []string) error { + sessionID := args[0] + mgr := getSessionManager() + ctx := context.Background() + + s, err := mgr.Get(ctx, sessionID) + if err != nil { + return fmt.Errorf("session not found: %s", sessionID) + } + + if s.IsLocked && !deleteForce { + return fmt.Errorf("session is locked by %s (use --force to override)", s.LockedBy) + } + + if s.IsLocked && deleteForce { + if err := mgr.ForceUnlock(ctx, sessionID); err != nil { + return fmt.Errorf("failed to force unlock: %w", err) + } + } + + if err := mgr.Delete(ctx, sessionID); err != nil { + return fmt.Errorf("failed to delete session: %w", err) + } + + fmt.Printf("Deleted session: %s\n", sessionID) + return nil +} + +func runSessionCleanup(cmd *cobra.Command, args []string) error { + duration, err := parseDuration(cleanupOlderThan) + if err != nil { + return fmt.Errorf("invalid duration: %w", err) + } + + mgr := getSessionManager() + ctx := context.Background() + + sessions, err := mgr.List(ctx) + if err != nil { + return fmt.Errorf("failed to list sessions: %w", err) + } + + cutoff := time.Now().Add(-duration) + var toDelete []*session.Session + + for _, s := range sessions { + if s.LastUsedAt.Before(cutoff) && !s.IsLocked { + toDelete = append(toDelete, s) + } + } + + if len(toDelete) == 0 { + fmt.Println("No sessions to clean up.") + return nil + } + + if cleanupDryRun { + fmt.Printf("Would delete %d session(s):\n", len(toDelete)) + for _, s := range toDelete { + size, _ := mgr.DiskUsage(s.ID) + fmt.Printf(" - %s (last used: %s, size: %s)\n", + s.ID, formatSessionDuration(time.Since(s.LastUsedAt)), formatSessionSize(size)) + } + return nil + } + + deleted, err := mgr.Cleanup(ctx, duration) + if err != nil { + return fmt.Errorf("cleanup failed: %w", err) + } + + fmt.Printf("Deleted %d session(s):\n", len(deleted)) + for _, id := range deleted { + fmt.Printf(" - %s\n", id) + } + + return nil +} + +func runSessionUnlock(cmd *cobra.Command, args []string) error { + sessionID := args[0] + mgr := getSessionManager() + ctx := context.Background() + + s, err := mgr.Get(ctx, sessionID) + if err != nil { + return fmt.Errorf("session not found: %s", sessionID) + } + + if !s.IsLocked { + fmt.Printf("Session %s is not locked.\n", sessionID) + return nil + } + + if err := mgr.ForceUnlock(ctx, sessionID); err != nil { + return fmt.Errorf("failed to unlock session: %w", err) + } + + fmt.Printf("Unlocked session: %s (was locked by %s)\n", sessionID, s.LockedBy) + return nil +} + +func formatSessionDuration(d time.Duration) string { + if d < time.Minute { + return fmt.Sprintf("%ds", int(d.Seconds())) + } + if d < time.Hour { + return fmt.Sprintf("%dm", int(d.Minutes())) + } + if d < 24*time.Hour { + return fmt.Sprintf("%dh", int(d.Hours())) + } + days := int(d.Hours() / 24) + return fmt.Sprintf("%dd", days) +} + +func formatSessionSize(bytes int64) string { + const ( + KB = 1024 + MB = KB * 1024 + GB = MB * 1024 + ) + switch { + case bytes >= GB: + return fmt.Sprintf("%.1f GB", float64(bytes)/GB) + case bytes >= MB: + return fmt.Sprintf("%.1f MB", float64(bytes)/MB) + case bytes >= KB: + return fmt.Sprintf("%.1f KB", float64(bytes)/KB) + default: + return fmt.Sprintf("%d B", bytes) + } +} + +func parseDuration(s string) (time.Duration, error) { + s = strings.TrimSpace(s) + if strings.HasSuffix(s, "d") { + days := strings.TrimSuffix(s, "d") + var d int + if _, err := fmt.Sscanf(days, "%d", &d); err != nil { + return 0, fmt.Errorf("invalid days: %s", days) + } + return time.Duration(d) * 24 * time.Hour, nil + } + return time.ParseDuration(s) +} diff --git a/docs/PRD_CONTAINERIZED_SANDBOX.md b/docs/PRD_CONTAINERIZED_SANDBOX.md index f4d0904d..826914df 100644 --- a/docs/PRD_CONTAINERIZED_SANDBOX.md +++ b/docs/PRD_CONTAINERIZED_SANDBOX.md @@ -501,6 +501,50 @@ For large files, use the NATS file staging system (`sandbox_stage_file` / `sandb ## Future Enhancements +### Private Registry Authentication + +For pulling images from private registries (ECR, GCR, Docker Hub private repos): + +**Configuration (config.yaml):** +```yaml +sandbox: + enabled: true + docker_image: "123456789.dkr.ecr.us-east-1.amazonaws.com/myapp:latest" + registry_auth: + # Option 1: Username/Password (Docker Hub, self-hosted) + username: "myuser" + password: "${DOCKER_TOKEN}" # Supports env var expansion + server_address: "https://index.docker.io/v1/" + + # Option 2: Identity Token (ECR, GCR, ACR) + identity_token: "${ECR_AUTH_TOKEN}" + server_address: "123456789.dkr.ecr.us-east-1.amazonaws.com" + + # Option 3: Docker config.json + docker_config_path: "/home/user/.docker/config.json" +``` + +**Environment Variables:** +| Variable | Description | +|----------|-------------| +| `STN_SANDBOX_DOCKER_IMAGE` | Docker image for sandbox containers | +| `STN_SANDBOX_REGISTRY_USERNAME` | Registry username | +| `STN_SANDBOX_REGISTRY_PASSWORD` | Registry password/token | +| `STN_SANDBOX_REGISTRY_TOKEN` | OAuth identity token (ECR/GCR) | +| `STN_SANDBOX_REGISTRY_SERVER` | Registry server address | +| `STN_SANDBOX_REGISTRY_CONFIG` | Path to Docker config.json | + +**CLI Configuration:** +```bash +stn config set sandbox.docker_image "ghcr.io/myorg/sandbox:latest" +stn config set sandbox.registry_auth.username "myuser" +stn config set sandbox.registry_auth.server_address "ghcr.io" +``` + +**UI:** Settings → Sandbox → Private Registry Auth section + +--- + ### Cloud Sandbox Fallback For environments without Docker (ECS Fargate, K8s, Fly.io): diff --git a/docs/agents/CREATING_AGENTS.md b/docs/agents/CREATING_AGENTS.md index 3979b9ca..2316d334 100644 --- a/docs/agents/CREATING_AGENTS.md +++ b/docs/agents/CREATING_AGENTS.md @@ -168,6 +168,52 @@ print(json.dumps(result)) See **[Sandbox Documentation](../station/sandbox.md)** for complete configuration options and examples. +### Agentic Harness Configuration (Advanced Multi-turn Execution) + +For agents that need Claude Agent SDK-like capabilities - multi-turn reasoning, persistent tool execution, and sandboxed environments - use the agentic harness: + +```yaml +harness: agentic +harness_config: + max_steps: 50 # Maximum reasoning steps (default: 25) + doom_loop_threshold: 5 # Abort after N similar outputs (default: 3) + timeout: 30m # Overall execution timeout (default: 15m) + sandbox: # Execution environment + mode: docker # host | docker | e2b (experimental) + image: python:3.11-slim # Docker image (if mode: docker) + network: false # Disable network access + memory: 4g # Memory limit +``` + +**Sandbox Modes:** +- `host` - Tools execute on host machine (fast, no isolation) +- `docker` - Tools execute in Docker containers (isolated, files persist via volume) +- `e2b` - Tools execute in cloud VMs (experimental, data doesn't persist) + +**Example Coding Agent:** +```yaml +--- +name: coding-agent +description: Multi-turn coding assistant +harness: agentic +harness_config: + max_steps: 50 + sandbox: + mode: docker + image: python:3.11 +tools: + - read + - write + - edit + - bash + - glob + - grep +--- +You are a coding assistant. Help users write, debug, and improve code. +``` + +See **[Harness Agents Documentation](../features/HARNESS_AGENTS.md)** for complete configuration options. + ### Tool Configuration ```yaml diff --git a/docs/features/HARNESS_AGENTS.md b/docs/features/HARNESS_AGENTS.md new file mode 100644 index 00000000..988eddde --- /dev/null +++ b/docs/features/HARNESS_AGENTS.md @@ -0,0 +1,376 @@ +# Agentic Harness - Advanced Agent Execution Engine + +The **Agentic Harness** is Station's Claude Agent SDK-like execution engine for agents that need full agentic capabilities: multi-turn reasoning, persistent tool execution, and isolated environments. + +## Overview + +The harness provides: +- **Multi-turn execution** with configurable step limits and doom loop protection +- **Sandboxed environments** for isolated tool execution (host, Docker, or E2B) +- **Built-in tools** for file operations, bash, and more +- **Streaming output** with real-time progress + +## Quick Start + +Create an agent with harness enabled: + +```yaml +--- +name: coding-agent +description: An agent that can write and modify code +harness: agentic +harness_config: + max_steps: 50 + doom_loop_threshold: 5 + timeout: 30m + sandbox: + mode: docker + image: python:3.11-slim +tools: + - read + - write + - bash + - glob + - grep + - edit +--- +You are a coding agent. Help users write, debug, and improve code. + +When given a task: +1. Understand the requirements +2. Plan your approach +3. Implement the solution +4. Verify it works + +Be thorough but efficient. +``` + +## Frontmatter Configuration + +### Basic Structure + +```yaml +--- +name: agent-name +description: What this agent does +harness: agentic # Enable agentic harness +harness_config: # Harness-specific configuration + max_steps: 50 # Max reasoning steps (default: 25) + doom_loop_threshold: 5 # Consecutive similar outputs before abort (default: 3) + timeout: 30m # Overall execution timeout (default: 15m) + sandbox: # Execution environment (optional) + mode: host # host | docker | e2b + # ... mode-specific options +tools: + - read + - write + - bash +--- +System prompt here... +``` + +### Harness Config Options + +| Option | Type | Default | Description | +|--------|------|---------|-------------| +| `max_steps` | int | 25 | Maximum agentic reasoning steps | +| `doom_loop_threshold` | int | 3 | Abort after N consecutive similar outputs | +| `timeout` | duration | 15m | Overall execution timeout | +| `sandbox` | object | null | Sandbox configuration (see below) | + +## Sandbox Modes + +The `harness_config.sandbox` determines WHERE tools execute: + +### Host Mode (Default) + +Tools execute directly on the host machine. Fast but no isolation. + +```yaml +harness_config: + sandbox: + mode: host +``` + +### Docker Mode + +Tools execute inside a Docker container with workspace volume mounting. + +```yaml +harness_config: + sandbox: + mode: docker + image: python:3.11-slim # Container image + network: false # Disable network access + timeout: 5m # Per-command timeout + memory: 4g # Memory limit + cpu: 2 # CPU limit + environment: # Environment variables + PYTHONPATH: /workspace +``` + +**Docker Options:** + +| Option | Type | Default | Description | +|--------|------|---------|-------------| +| `image` | string | python:3.11-slim | Docker image | +| `network` | bool | false | Enable network access | +| `timeout` | duration | 5m | Per-command timeout | +| `memory` | string | 4g | Memory limit (Docker format) | +| `cpu` | int | 2 | CPU core limit | +| `environment` | map | {} | Environment variables | + +**How Docker Mode Works:** +- Workspace is mounted as a volume (`-v /host/path:/workspace`) +- Files persist across container destroys +- Each tool execution gets a fresh container but sees same files +- Great for reproducible, isolated execution + +### E2B Mode (Experimental) + +Tools execute in E2B cloud VMs. + +```yaml +harness_config: + sandbox: + mode: e2b + template: base # E2B template + timeout: 10m # VM timeout +``` + +> **Note:** E2B mode is experimental. Data does not persist across sandbox destroys - the harness would need CopyIn/CopyOut sync between agent steps for stateful workflows. Use Docker mode for production workloads. + +## Built-in Tools + +The harness provides these tools by default when enabled: + +| Tool | Description | +|------|-------------| +| `read` | Read file contents | +| `write` | Write/create files | +| `edit` | Edit files with search/replace | +| `bash` | Execute shell commands | +| `glob` | Find files by pattern | +| `grep` | Search file contents | + +Add them to your agent's `tools:` list: + +```yaml +tools: + - read + - write + - bash + - glob + - grep + - edit +``` + +## Harness vs Legacy Sandbox + +**Important:** There are two sandbox concepts in Station: + +| Feature | Legacy `sandbox:` | Harness `harness_config.sandbox:` | +|---------|-------------------|-----------------------------------| +| Purpose | One-shot code execution via `sandbox_run` tool | Persistent environment for agentic tools | +| Location | Root frontmatter | Nested under `harness_config` | +| Runtime | Dagger-based | Docker or E2B | +| Use Case | Single script execution | Multi-turn agent workflows | + +**Legacy sandbox (still works for non-harness agents):** +```yaml +--- +name: simple-agent +sandbox: # OLD - for sandbox_run tool + mode: compute + runtime: python +tools: + - sandbox_run +--- +``` + +**Harness sandbox:** +```yaml +--- +name: coding-agent +harness: agentic +harness_config: + sandbox: # NEW - execution environment + mode: docker +tools: + - read + - write + - bash +--- +``` + +## Example Agents + +### Python Development Agent + +```yaml +--- +name: python-dev +description: Python development assistant +harness: agentic +harness_config: + max_steps: 100 + timeout: 1h + sandbox: + mode: docker + image: python:3.11 + environment: + PIP_NO_CACHE_DIR: "1" +tools: + - read + - write + - edit + - bash + - glob + - grep +--- +You are a Python development assistant. + +You can: +- Write and modify Python code +- Run tests with pytest +- Install packages with pip +- Debug issues + +Always: +1. Read existing code before modifying +2. Run tests after changes +3. Explain what you changed and why +``` + +### Infrastructure Agent + +```yaml +--- +name: infra-agent +description: Infrastructure and DevOps tasks +harness: agentic +harness_config: + max_steps: 50 + sandbox: + mode: docker + image: hashicorp/terraform:latest + network: true # Need network for cloud APIs +tools: + - read + - write + - bash + - glob +--- +You are an infrastructure assistant specializing in Terraform and cloud resources. + +Help users: +- Write Terraform configurations +- Plan and validate changes +- Debug infrastructure issues +``` + +## Creating Harness Agents + +### Via MCP (create_agent tool) + +```json +{ + "name": "my-harness-agent", + "description": "Agent with agentic harness", + "prompt": "You are a helpful agent...", + "environment_id": "1", + "harness_config": { + "max_steps": 50, + "doom_loop_threshold": 5, + "timeout": "30m", + "sandbox": { + "mode": "docker", + "image": "python:3.11" + } + }, + "tool_names": ["read", "write", "bash", "glob", "grep", "edit"] +} +``` + +### Via CLI + +```bash +stn agent create my-agent \ + --prompt "You are a helpful agent..." \ + --description "Agent with harness" \ + --harness-config '{"max_steps": 50, "sandbox": {"mode": "docker"}}' \ + --tools read,write,bash,glob,grep,edit +``` + +### Via .prompt File + +Create `environments/default/agents/my-agent.prompt`: + +```yaml +--- +name: my-agent +description: Agent with harness +harness: agentic +harness_config: + max_steps: 50 + sandbox: + mode: docker +tools: + - read + - write + - bash +--- +Your system prompt here... +``` + +Then sync the environment: + +```bash +stn sync default +``` + +## Execution Flow + +1. **Task Received** - Agent receives user task +2. **Sandbox Created** - If configured, sandbox environment is provisioned +3. **Tool Registry Built** - Tools are registered with sandbox executor +4. **Agentic Loop** - Agent reasons and calls tools iteratively +5. **Doom Loop Check** - Detects repetitive behavior +6. **Step Limit Check** - Enforces max_steps +7. **Timeout Check** - Enforces overall timeout +8. **Result Returned** - Final response or error + +## Troubleshooting + +### "Doom loop detected" + +Agent is producing repetitive outputs. Solutions: +- Improve system prompt with clearer instructions +- Increase `doom_loop_threshold` if behavior is legitimate +- Check if the task is achievable + +### Docker container fails to start + +- Ensure Docker is running: `docker ps` +- Check image exists: `docker pull ` +- Verify memory/CPU limits are reasonable + +### E2B timeout + +- E2B VMs have session limits +- Consider Docker mode for longer tasks +- Check E2B API key is valid + +### Tools not found + +- Ensure tools are listed in `tools:` array +- Harness tools (read, write, bash, etc.) must be explicitly listed +- MCP tools require the MCP server to be running + +## Best Practices + +1. **Start with host mode** for development, switch to Docker for production +2. **Set reasonable timeouts** - longer isn't always better +3. **Use doom loop protection** - prevents runaway agents +4. **List only needed tools** - fewer tools = better focus +5. **Test in isolation** - use Docker for reproducible results diff --git a/docs/features/HARNESS_ARCHITECTURE_ANALYSIS.md b/docs/features/HARNESS_ARCHITECTURE_ANALYSIS.md new file mode 100644 index 00000000..1d7a5d9d --- /dev/null +++ b/docs/features/HARNESS_ARCHITECTURE_ANALYSIS.md @@ -0,0 +1,406 @@ +# Harness Architecture Analysis + +**Date**: 2025-01-07 +**Purpose**: Understand how sandbox integration fits into the broader harness ecosystem + +## Executive Summary + +The sandbox work in `pkg/harness/sandbox/` is **architecturally isolated** and **not wired into the actual execution path**. While we built comprehensive sandbox interfaces (Host, Docker, E2B), the core execution engine (`executeWithAgenticHarness`) doesn't use them. + +### Critical Gaps Identified + +| Gap | Impact | Priority | +|-----|--------|----------| +| `executeWithAgenticHarness` doesn't use sandbox | Tools execute directly on host | **HIGH** | +| ToolRegistry created without sandbox | All tools bypass sandbox isolation | **HIGH** | +| No sandbox config in harness frontmatter | Can't enable per-agent | **MEDIUM** | +| Session manager exists but unused | Workspace isolation incomplete | **MEDIUM** | +| Existing `SandboxService` in `internal/services/` | Potential duplication | **LOW** | + +--- + +## Current Architecture Overview + +``` +┌─────────────────────────────────────────────────────────────────────────────┐ +│ STATION ARCHITECTURE │ +├─────────────────────────────────────────────────────────────────────────────┤ +│ │ +│ Entry Points: │ +│ ├── CLI: stn call │ +│ ├── API: /api/agents/:id/execute │ +│ ├── Workflow Engine: internal/workflows/runtime/ │ +│ └── Lattice: ExecutorAdapter (remote execution) │ +│ │ +│ │ │ +│ ▼ │ +│ ┌───────────────────────────────────────────────────────────────────────┐ │ +│ │ AgentExecutionEngine.ExecuteWithOptions() │ │ +│ │ │ │ +│ │ 1. Load harness frontmatter from .prompt file │ │ +│ │ 2. Check: if harnessMode == "agentic" → executeWithAgenticHarness() │ │ +│ │ 3. Otherwise → GenKitExecutor.ExecuteAgent() (dotprompt default) │ │ +│ └───────────────────────────────────────────────────────────────────────┘ │ +│ │ │ +│ ┌───────────────┴───────────────┐ │ +│ ▼ ▼ │ +│ ┌─────────────────────────┐ ┌─────────────────────────┐ │ +│ │ Agentic Harness │ │ Dotprompt (Default) │ │ +│ │ (pkg/harness/) │ │ (pkg/dotprompt/) │ │ +│ │ │ │ │ │ +│ │ • Manual loop control │ │ • Genkit's default │ │ +│ │ • Doom loop detection │ │ • Simple execution │ │ +│ │ • Context compaction │ │ • No special features │ │ +│ │ • Git integration │ │ │ │ +│ │ • Streaming events │ │ │ │ +│ └─────────────────────────┘ └─────────────────────────┘ │ +│ │ +└─────────────────────────────────────────────────────────────────────────────┘ +``` + +--- + +## The Problem: Sandbox Not Wired + +### Current `executeWithAgenticHarness()` (lines 1391-1512) + +```go +func (aee *AgentExecutionEngine) executeWithAgenticHarness(...) { + // 1. Load config + harnessConfig := convertConfigToHarnessConfig(loadedConfig.Harness) + + // 2. Create workspace (host only!) + ws := workspace.NewHostWorkspace(workspacePath) // ❌ No sandbox option + + // 3. Create tool registry WITHOUT sandbox + toolRegistry := tools.NewToolRegistry(genkitApp, workspacePath) // ❌ Missing sandbox! + toolRegistry.RegisterBuiltinTools() + + // 4. Create executor (has sandbox fields but never used) + executorOpts := []harness.ExecutorOption{ + harness.WithWorkspace(ws), + harness.WithModelName(modelName), + // ❌ NO WithSandbox() or WithSandboxConfig()! + } + + agenticExecutor := harness.NewAgenticExecutor(...) + // agenticExecutor.sandbox is nil + // Tools execute directly on host +} +``` + +### What Should Happen + +```go +func (aee *AgentExecutionEngine) executeWithAgenticHarness(...) { + // 1. Parse sandbox config from frontmatter + sandboxConfig := parseSandboxConfig(harnessFM) // NEW + + // 2. Create RuntimeResolver + resolver := sandbox.NewRuntimeResolver(sandbox.IsolationDocker) // or E2B + + // 3. Create sandbox based on config + sb, err := createSandbox(ctx, sandboxConfig, resolver) // NEW + + // 4. Create tool registry WITH sandbox + toolRegistry := tools.NewToolRegistryWithSandbox(genkitApp, workspacePath, sb) // ✅ + + // 5. Create executor WITH sandbox + executorOpts := []harness.ExecutorOption{ + harness.WithSandbox(sb), // ✅ + harness.WithSandboxConfig(sandboxConfig), // ✅ + harness.WithRuntimeResolver(resolver), // ✅ + } +} +``` + +--- + +## Component Inventory + +### What We Built (pkg/harness/) + +| Component | Location | Status | Purpose | +|-----------|----------|--------|---------| +| **Sandbox Interface** | `sandbox/sandbox.go` | ✅ Complete | Abstract sandbox interface | +| **HostSandbox** | `sandbox/host.go` | ✅ Complete | Passthrough (no isolation) | +| **DockerSandbox** | `sandbox/docker.go` | ✅ Complete | Container isolation | +| **E2BSandbox** | `sandbox/e2b.go` | ✅ Complete | Cloud sandbox | +| **StreamingSandbox** | `sandbox/sandbox.go` | ✅ Complete | Real-time output | +| **RuntimeResolver** | `sandbox/sandbox.go` | ✅ Complete | Auto-detect mode | +| **Factory** | `sandbox/sandbox.go` | ✅ Complete | Create sandboxes | +| **ToolRegistry w/Sandbox** | `tools/registry.go` | ✅ Complete | Tools use sandbox | +| **Bash Tool w/Sandbox** | `tools/bash.go` | ✅ Complete | Executes via sandbox | +| **AgenticExecutor Sandbox Fields** | `executor.go` | ✅ Complete | Fields exist | +| **Executor Sandbox Lifecycle** | `executor.go` | ✅ Complete | setup/cleanup wired | + +### What's Missing (Integration) + +| Component | Location | Status | Required Action | +|-----------|----------|--------|-----------------| +| **Sandbox frontmatter parsing** | `agent_execution_engine.go` | ❌ Missing | Parse `sandbox:` from .prompt | +| **Sandbox creation in engine** | `agent_execution_engine.go` | ❌ Missing | Create sandbox before tools | +| **Pass sandbox to ToolRegistry** | `agent_execution_engine.go` | ❌ Missing | Use `NewToolRegistryWithSandbox` | +| **Pass sandbox to executor** | `agent_execution_engine.go` | ❌ Missing | Use `WithSandbox()` option | +| **Session → Workspace → Sandbox** | `session/manager.go` | ❌ Disconnected | Wire session to sandbox | + +--- + +## Existing Station Sandbox System + +Station already has a sandbox system in `internal/services/`: + +``` +internal/services/ +├── sandbox_service.go # Core service (Dagger-based) +├── sandbox_docker_backend.go # Docker backend for persistent sessions +└── sandbox_opencode_backend.go # OpenCode integration +``` + +### Comparison + +| Feature | Existing (`internal/services/`) | New (`pkg/harness/sandbox/`) | +|---------|--------------------------------|------------------------------| +| **Engine** | Dagger | Direct Docker/E2B API | +| **Purpose** | Code execution tool | Agent execution isolation | +| **Scope** | Single code snippet | Entire agent run | +| **Streaming** | No | Yes (ExecStream) | +| **Auto-detect** | No | Yes (RuntimeResolver) | +| **E2B Support** | No | Yes | + +### Recommendation + +The `pkg/harness/sandbox/` is better suited for agent-level isolation because: +1. Designed for entire agent runs, not single code executions +2. Streaming support for real-time tool output +3. RuntimeResolver for environment adaptation +4. E2B support for serverless deployment + +--- + +## Deployment Constraint Analysis + +### Why Sandbox Was Needed + +Station runs in containers via `stn up`: + +```dockerfile +# Dockerfile - Station container +FROM ubuntu:22.04 +# Contains: git, python, node, docker CLI +# But NOT: arbitrary dev tools the agent might need +``` + +**Problem**: When Station runs in a container, agents have limited tools: +- Container has basic tools (git, python, node) +- Agent might need tools not in container (rust, terraform, kubectl) +- Can't install tools at runtime (container is immutable) + +**Solution Options**: + +| Option | How It Works | Pros | Cons | +|--------|--------------|------|------| +| **Host Mode** | Execute on host | Full access | No isolation | +| **Docker Mode** | Spawn sibling containers | Custom images | Docker-in-Docker complexity | +| **E2B Mode** | Cloud sandbox | No local Docker needed | Latency, cost | + +### RuntimeResolver Flow + +``` +┌─────────────────────────────────────────────────────────────────┐ +│ DEPLOYMENT SCENARIO │ +├─────────────────────────────────────────────────────────────────┤ +│ │ +│ SCENARIO 1: Developer laptop │ +│ ├── Host has all tools (python, node, etc.) │ +│ └── RuntimeResolver → ModeHost (no overhead) │ +│ │ +│ SCENARIO 2: Station container (stn up) │ +│ ├── Container has limited tools │ +│ ├── Agent needs tools not in container │ +│ ├── isolation_backend: docker configured │ +│ └── RuntimeResolver → ModeDocker (spawn sibling container) │ +│ │ +│ SCENARIO 3: Serverless/CloudShip deployment │ +│ ├── No Docker available │ +│ ├── isolation_backend: e2b configured │ +│ └── RuntimeResolver → ModeE2B (cloud sandbox) │ +│ │ +└─────────────────────────────────────────────────────────────────┘ +``` + +--- + +## Session / Workspace / Sandbox Relationship + +### Current State (Disconnected) + +``` +Session (pkg/harness/session/) +├── Manages: workspace paths, locks, metadata +├── NOT connected to: sandbox +└── Used by: nothing currently + +Workspace (pkg/harness/workspace/) +├── Manages: file operations within a path +├── NOT connected to: sandbox +└── Used by: executeWithAgenticHarness (HostWorkspace only) + +Sandbox (pkg/harness/sandbox/) +├── Manages: isolated execution environment +├── NOT connected to: session or workspace +└── Used by: nothing currently +``` + +### Desired State (Connected) + +``` +Session +├── Creates/manages workspace path +├── Provides repo URL for cloning +└── Locks workspace during execution + │ + ▼ +Workspace Resolver +├── Resolves workspace path from session +├── Clones repo if git_source provided +└── Returns workspace path + │ + ▼ +Sandbox +├── Receives workspace path as mount point +├── Mounts workspace into container (Docker mode) +├── Uploads workspace to cloud (E2B mode) +└── Tools execute inside sandbox with workspace access + │ + ▼ +Tools +├── Read/write files via sandbox.ReadFile/WriteFile +├── Execute commands via sandbox.Exec +└── Changes visible in workspace after completion +``` + +--- + +## NATS / Lattice Integration + +### Current Integration Points + +``` +pkg/harness/nats/ +├── lattice_adapter.go # Uses existing work.WorkStore +├── store.go # NATS KV + Object Store +├── handoff.go # Multi-agent workflow coordination +└── nats_publisher.go # Streaming events to NATS + +pkg/harness/stream/ +└── nats_publisher.go # Event streaming +``` + +### How Sandbox Affects Lattice + +Sandbox doesn't change lattice architecture, but: +1. **Work records** should track sandbox mode used +2. **File sharing** may need to extract from sandbox to Object Store +3. **Streaming** should include sandbox events (container start/stop) + +--- + +## Integration Roadmap + +### Phase 3.5: Wire Sandbox into Execution (HIGH PRIORITY) + +``` +1. Modify executeWithAgenticHarness(): + - Parse sandbox config from frontmatter + - Create sandbox via Factory + - Pass to ToolRegistry + - Pass to AgenticExecutor + +2. Add frontmatter support: + sandbox: + mode: auto | host | docker | e2b + image: station-sandbox:latest # For docker mode + timeout: 30m + +3. Add config.yaml support: + harness: + sandbox: + default_mode: auto + isolation_backend: docker # Fallback for auto mode + e2b_api_key: ${E2B_API_KEY} +``` + +### Phase 3.6: Session → Workspace → Sandbox (MEDIUM PRIORITY) + +``` +1. Wire SessionManager into executeWithAgenticHarness() +2. Resolve workspace from session +3. Clone repo if git_source provided +4. Create sandbox with workspace mounted +5. Lock session during execution +``` + +### Phase 3.7: NATS Sandbox Events (LOW PRIORITY) + +``` +1. Add sandbox_start/sandbox_stop events to stream +2. Track sandbox mode in work records +3. Extract artifacts from sandbox to Object Store +``` + +--- + +## Test Status + +### Tests That Pass (Isolated) + +```bash +go test ./pkg/harness/sandbox/... -v # 25 tests ✅ +go test ./pkg/harness/... -v # All pass ✅ +``` + +### Tests That Would Fail (If Integration Existed) + +The sandbox tests pass because they test the sandbox in isolation. If we wire it into `executeWithAgenticHarness()`, we'd need: + +1. **E2E tests with sandbox** - Currently skip sandbox entirely +2. **Docker-in-Docker tests** - Need DinD test infrastructure +3. **E2B integration tests** - Need E2B API key and network access + +--- + +## Recommendations + +### Immediate (This Session) + +1. **Don't** try to fully wire sandbox into agent_execution_engine.go +2. **Do** document the integration plan clearly +3. **Do** ensure pkg/harness/sandbox tests continue to pass + +### Next Session + +1. Create `parseHarnessSandboxConfig()` function +2. Add sandbox frontmatter schema +3. Wire sandbox into `executeWithAgenticHarness()` +4. Add integration test with Docker sandbox + +### Future + +1. Wire session manager +2. Add E2B integration test +3. Add sandbox events to NATS stream + +--- + +## Files Reference + +| File | Purpose | Lines | +|------|---------|-------| +| `internal/services/agent_execution_engine.go` | Main entry, needs sandbox wiring | 1391-1512 | +| `pkg/harness/executor.go` | Has sandbox fields, lifecycle wired | ~600 | +| `pkg/harness/tools/registry.go` | Has `NewToolRegistryWithSandbox` | ~90 | +| `pkg/harness/tools/bash.go` | Uses sandbox if provided | ~190 | +| `pkg/harness/sandbox/sandbox.go` | Interface + Factory + Resolver | ~310 | +| `pkg/harness/session/manager.go` | Session management (unused) | ~360 | +| `pkg/harness/nats/lattice_adapter.go` | NATS integration | ~530 | diff --git a/docs/features/PRD_AGENTIC_HARNESS.md b/docs/features/PRD_AGENTIC_HARNESS.md new file mode 100644 index 00000000..c46239c3 --- /dev/null +++ b/docs/features/PRD_AGENTIC_HARNESS.md @@ -0,0 +1,1248 @@ +# PRD: Agentic Harness - Claude Agent SDK-like Execution Engine + +**Status**: Phase 6 Complete (Interactive REPL with Session Persistence) +**Author**: Claude/Human Collaboration +**Created**: 2025-01-06 +**Last Updated**: 2025-01-26 + +## Overview + +The Agentic Harness is an alternative execution engine for Station agents that provides Claude Agent SDK-like capabilities: manual agentic loop control, doom loop detection, context compaction, git integration, and workspace management. + +## Problem Statement + +The default Genkit-based execution (`dotprompt.ExecuteAgent`) works well for most agents but has limitations: + +1. **No doom loop detection** - Agents can get stuck in repetitive patterns +2. **No context compaction** - Long conversations can exceed context windows +3. **Limited git integration** - No automatic branch/commit management +4. **No workspace isolation** - Agents share the same file system context +5. **No manual loop control** - Cannot customize step behavior or hooks + +## Solution + +Add `harness: agentic` support to agent dotprompt files, enabling a Claude Agent SDK-like execution model with: + +- Manual agentic loop with step-by-step control +- Doom loop detection and prevention +- Context compaction for long-running tasks +- Git integration (auto-branch, auto-commit) +- Workspace management and isolation +- Pre/post tool execution hooks + +## Architecture + +### Integration Point + +The harness integrates at the `AgentExecutionEngine.Execute()` level: + +``` +┌─────────────────────────────────────────────────────────────────────┐ +│ AgentExecutionEngine.Execute() │ +│ │ │ +│ ┌────────────────────┴────────────────────┐ │ +│ ▼ ▼ │ +│ harness: "" (default) harness: agentic │ +│ │ │ │ +│ ▼ ▼ │ +│ dotprompt.ExecuteAgent() executeWithAgenticHarness() │ +│ (Genkit-based execution) (Manual agentic loop) │ +│ │ │ +│ ▼ │ +│ pkg/harness/AgenticExecutor │ +│ • Doom loop detection │ +│ • Context compaction │ +│ • Git integration │ +│ • Workspace management │ +└─────────────────────────────────────────────────────────────────────┘ +``` + +### Component Structure + +``` +pkg/harness/ +├── executor.go # AgenticExecutor - main execution engine +├── config.go # HarnessConfig, AgentHarnessConfig +├── doom_loop.go # DoomLoopDetector +├── compaction.go # Compactor for context management +├── hooks.go # HookRegistry for pre/post tool hooks +├── tools/ +│ ├── registry.go # ToolRegistry +│ ├── read.go # File read tool +│ ├── write.go # File write tool +│ ├── edit.go # File edit tool +│ ├── bash.go # Bash command tool +│ ├── glob.go # File glob tool +│ ├── grep.go # Content search tool +│ └── git_tools.go # Git operation tools +├── workspace/ +│ └── host.go # HostWorkspace implementation +├── git/ +│ └── manager.go # GitManager for branch/commit operations +├── prompt/ +│ └── builder.go # PromptBuilder for system prompts +└── trace/ + └── tracer.go # OpenTelemetry tracing integration +``` + +## Usage + +### Agent Configuration + +Add `harness: agentic` to the agent's dotprompt frontmatter: + +```yaml +--- +model: anthropic/claude-sonnet-4-20250514 +harness: agentic +harness_config: + max_steps: 50 + doom_loop_threshold: 3 + timeout: 10m +tools: + - read_file + - write_file + - bash +--- +You are a code analysis agent... +``` + +### Global Configuration + +Configure harness defaults in `config.yaml`. Running `stn init` sets these defaults automatically: + +```yaml +harness: + workspace: + path: ./workspace + mode: host # or "sandbox" + compaction: + enabled: true + threshold: 0.85 # Compact at 85% of context window + protect_tokens: 40000 # Keep last N tokens from compaction + git: + auto_branch: true + branch_prefix: agent/ + auto_commit: false + require_approval: true + workflow_branch_strategy: shared + nats: + enabled: true + kv_bucket: harness-state + object_bucket: harness-files + max_file_size: 100MB + ttl: 24h + permissions: + external_directory: deny + # bash permissions use code defaults (see DefaultHarnessConfig) +``` + +The `stn init` command also creates the workspace directory at `~/.config/station/workspace/`. + +## Key Features + +### 1. Doom Loop Detection + +Detects when an agent is stuck in a repetitive pattern: + +```go +type DoomLoopDetector struct { + threshold int // Number of similar calls to trigger + history []ToolCallRecord // Recent tool calls + maxHistory int // Maximum history to keep +} + +// Triggers when same tool+input seen N times consecutively +func (d *DoomLoopDetector) IsInLoop() bool +``` + +**Configuration**: +```yaml +harness_config: + doom_loop_threshold: 3 # Trigger after 3 identical calls +``` + +### 2. Context Compaction + +Automatically summarizes conversation history when approaching context limits: + +```go +type Compactor struct { + genkitApp *genkit.Genkit + modelName string + config CompactionConfig + contextWindow int +} + +// Returns compacted history when token count exceeds threshold +func (c *Compactor) CompactIfNeeded(ctx context.Context, history []*ai.Message) ([]*ai.Message, bool, error) +``` + +**Configuration**: +```yaml +harness: + compaction: + enabled: true + threshold: 0.8 # 80% of context window + protect_tokens: 2000 # Protect recent tokens +``` + +### 3. Git Integration + +Automatic branch creation and commits: + +```go +type GitManager interface { + CreateBranch(ctx context.Context, task string, agentID string) (string, error) + Commit(ctx context.Context, message string) (string, error) + Push(ctx context.Context) error + GetCurrentBranch(ctx context.Context) (string, error) +} +``` + +**Configuration**: +```yaml +harness: + git: + auto_branch: true + branch_prefix: agent/ + auto_commit: true +``` + +### 4. Built-in Tools + +The harness provides built-in tools that work independently of MCP: + +| Tool | Description | +|------|-------------| +| `read` | Read file contents | +| `write` | Write file contents | +| `edit` | Edit file with string replacement | +| `bash` | Execute bash commands | +| `glob` | Find files by pattern | +| `grep` | Search file contents | +| `git_status` | Get git status | +| `git_diff` | Get git diff | +| `git_log` | Get git log | + +### 5. Hook System + +Pre/post tool execution hooks for: +- Permission checks +- Audit logging +- Doom loop detection +- Custom validation + +```go +type HookResult string +const ( + HookContinue HookResult = "continue" + HookBlock HookResult = "block" + HookInterrupt HookResult = "interrupt" +) + +type PreToolHook func(ctx context.Context, req *ai.ToolRequest) (HookResult, string) +type PostToolHook func(ctx context.Context, req *ai.ToolRequest, result interface{}) +``` + +## Workflow Integration + +The harness works with Station's workflow system. Agents with `harness: agentic` are executed through `executeWithAgenticHarness()` when triggered by workflows: + +```yaml +# workflow.yaml +id: code-review +states: + - name: analyze + type: agent + agent: code-analyzer # Has harness: agentic in dotprompt + transition: report +``` + +## Multi-Agent Handoff + +The `pkg/harness/nats/LatticeAdapter` bridges harness agents with Station's existing lattice infrastructure: + +```go +// LatticeAdapter uses existing WorkStore for state persistence +adapter, err := nats.NewLatticeAdapter(js, nats.LatticeAdapterConfig{ + StationID: stationID, + WorkStoreConfig: work.DefaultWorkStoreConfig(), + FileStoreConfig: nats.DefaultFileStoreConfig(), +}) + +// Create work record for harness execution +record, err := adapter.CreateHarnessWork(ctx, nats.CreateHarnessWorkInput{ + WorkID: workID, + WorkflowRunID: workflowRunID, + StepID: stepID, + AgentID: agentID, + AgentName: agentName, + Task: task, + OrchestratorRunID: orchestratorRunID, +}) + +// Upload/download files between workflow steps +adapter.UploadOutputFile(ctx, workID, localPath) +adapter.DownloadPreviousStepFiles(ctx, orchestratorRunID, localDir) +``` + +The adapter stores workflow metadata (WorkflowRunID, StepID, GitBranch) in the `Context` map of `work.WorkRecord`, maintaining compatibility with the existing lattice infrastructure. + +## Testing + +### Unit Tests +```bash +go test ./pkg/harness/... -v +``` + +### Integration Tests (with real LLM) +```bash +HARNESS_E2E_TEST=1 go test ./pkg/harness/... -tags=integration -run TestAgenticExecutor -v +``` + +### Multi-Agent E2E Test +```bash +HARNESS_E2E_TEST=1 go test ./pkg/harness/... -tags=integration -run TestMultiAgentWorkflow -v -timeout 10m +``` + +## Success Criteria + +- [x] Harness execution mode selectable via `harness: agentic` frontmatter +- [x] Doom loop detection prevents infinite loops +- [x] Built-in tools work independently of MCP +- [x] Git integration creates branches and commits +- [x] Workspace isolation prevents file conflicts +- [x] Context compaction tested with real LLM (E2E test with 1000 token window) +- [x] Multi-agent handoff uses existing lattice infrastructure (LatticeAdapter in pkg/harness/nats/) +- [x] E2E test with workflow simulation passes (TestAgenticExecutor_E2E_WorkflowSimulation) +- [x] `stn init` provides sensible harness defaults + +## Risks & Mitigations + +| Risk | Mitigation | +|------|------------| +| Doom loop false positives | Configurable threshold, tool+input hashing | +| Compaction loses context | Protected token window, summarization prompt | +| Git conflicts | Branch per agent/task, require approval option | +| MCP tool conflicts | Built-in tools have distinct names | + +## Timeline + +| Phase | Status | Description | +|-------|--------|-------------| +| Phase 1-8 | Complete | Core harness, tools, doom loop, workspace, git | +| Phase 9 | Complete | NATS integration via LatticeAdapter (uses work.WorkStore) | +| Phase 10 | Complete | Multi-agent E2E test with real LLM | +| Phase 11 | Complete | Compaction E2E test with configurable token window | +| Phase 12 | Complete | Workflow simulation E2E test (multi-agent pipeline) | + +## E2E Test Results (2025-01-06) + +All tests pass with Claude Opus 4.5: + +| Test | Duration | Tokens | Description | +|------|----------|--------|-------------| +| `TestAgenticExecutor_E2E_RealLLM` | ~7s | ~2,100 | Basic file create/read | +| `TestAgenticExecutor_E2E_MultiStep` | ~11s | ~2,700 | Directory + JSON config | +| `TestAgenticExecutor_E2E_Compaction` | ~10s | ~2,500 | Compaction config wired | +| `TestAgenticExecutor_E2E_WorkflowSimulation` | ~13s | ~4,400 | Multi-agent workflow | + +Run tests: +```bash +HARNESS_E2E_TEST=1 go test ./pkg/harness/... -v -run "E2E" -timeout 5m +``` + +## Phase 2: Workspace Isolation & Session Persistence (IN PROGRESS) + +### Problem Statement + +The current harness uses a single shared workspace path for all executions. This causes issues: + +1. **Concurrent runs conflict** - Two agents writing to same files +2. **No continuation** - Can't resume work from previous run +3. **No repo targeting** - Can't tell agent "work on this specific repo" +4. **Workflow vs standalone** - Different isolation needs + +### Invocation Sources + +Agents can be invoked from multiple sources, each with different workspace needs: + +| Source | Identity Provider | Example | +|--------|------------------|---------| +| CLI | User provides `--session` or `--repo` | `stn call coder --session proj-123` | +| Workflow | Engine provides `workflow_run_id` | Steps share workspace automatically | +| Agent handoff | Parent context propagates | Child inherits parent's session | +| Schedule | Config specifies | `schedule.session_id: "nightly-scan"` | +| Event/webhook | Payload contains | `{"session_id": "pr-456", "repo": "..."}` | +| CloudShip | Remote orchestration | Dispatch includes repo URL | + +### Proposed Design + +``` +┌─────────────────────────────────────────────────────────────────┐ +│ RUNTIME PARAMETERS │ +│ (passed at invocation time, not config) │ +├─────────────────────────────────────────────────────────────────┤ +│ session_id - User-provided, persists across runs │ +│ workflow_run_id - From workflow engine context │ +│ agent_run_id - Auto-generated per execution │ +│ git_source - {url, branch, ref} to clone │ +└─────────────────────────────────────────────────────────────────┘ + │ + ▼ +┌─────────────────────────────────────────────────────────────────┐ +│ WORKSPACE RESOLVER │ +│ │ +│ Resolution Priority: │ +│ 1. session_id provided → workspace/session/{session_id}/ │ +│ 2. workflow_run_id present → workspace/workflow/{wf_run_id}/ │ +│ 3. fallback → workspace/run/{agent_run_id}/ │ +│ │ +│ Git Source Handling: │ +│ - If git_source provided AND workspace empty → clone repo │ +│ - If git_source provided AND workspace exists → fetch/checkout │ +│ - If no git_source → use workspace as-is │ +└─────────────────────────────────────────────────────────────────┘ +``` + +### Configuration + +**Global defaults (config.yaml):** +```yaml +harness: + workspace: + path: ./workspace + mode: host + isolation_mode: per_workflow # shared | per_run | per_workflow + cleanup_on_complete: false + cleanup_after: 24h +``` + +**Runtime parameters (passed at invocation):** +```go +type ExecutionContext struct { + SessionID string // User-provided, highest priority + WorkflowRunID string // From workflow context + AgentRunID string // Auto-generated + AgentName string // For branch naming + Task string // For branch naming + GitSource *GitSourceConfig // Optional repo to clone +} +``` + +**CLI usage:** +```bash +# Fresh isolated run (default) +stn call my-agent --task "analyze code" + +# Continue from session +stn call my-agent --task "continue refactoring" --session "refactor-auth" + +# Work on specific repo +stn call my-agent --task "review PR" --repo "https://github.com/org/repo" --ref "pr-123" + +# Combine: specific repo with persistent session +stn call my-agent --task "fix tests" --repo "git@github.com:org/repo" --session "fix-tests-sprint-42" +``` + +**Workflow usage:** +```yaml +id: code-review-pipeline +input: + repo_url: string + session_id: string # Optional: persist across workflow runs + +states: + - name: clone + type: agent + agent: repo-setup + input: + git_source: + url: "${repo_url}" + depth: 1 + session_id: "${session_id}" # All steps share this + + - name: analyze + type: agent + agent: code-analyzer + # Inherits session from workflow context +``` + +### Design Decisions + +#### 1. Collaboration Modes + +Two modes for how agents interact within a session: + +``` +SEQUENTIAL MODE (default) +━━━━━━━━━━━━━━━━━━━━━━━━━ +Agent A ──commits──► Agent B ──commits──► Agent C + same branch (pick up where left off) + +Use: Pipeline where each step builds on previous +Example: analyze → fix → test → document + + +PARALLEL MODE (explicit) +━━━━━━━━━━━━━━━━━━━━━━━━━ + ┌── Agent A (branch: agent/session/frontend) +base ────┼── Agent B (branch: agent/session/backend) + └── Agent C (branch: agent/session/docs) + │ + ▼ + merge step + +Use: Independent work that gets combined later +Example: parallel feature development, multiple reviewers +``` + +**Configuration:** +```yaml +harness: + git: + collaboration: sequential # sequential | parallel + working_branch: "" # Empty = current branch +``` + +#### 2. Concurrent Access + +**Decision: Mode-dependent behavior** + +| Mode | Same session, same time | Behavior | +|------|------------------------|----------| +| Sequential | Second caller blocked | "Session locked by run X" | +| Parallel | Allowed | Each gets own branch | + +For sequential mode, we use advisory file locks on the session directory. + +#### 3. Session Lifecycle + +**Decision: Explicit management with warnings** + +```bash +stn session list # Shows sessions + last-used time +stn session delete proj-123 # Explicit cleanup +stn session cleanup --older-than 7d # Bulk cleanup +``` + +Sessions persist until explicitly deleted. Station warns about stale sessions. + +#### 4. Workflow Configuration + +**Workflow-level git settings:** +```yaml +id: code-review-pipeline +git: + collaboration: sequential # Default for all steps + working_branch: feature/review-${input.pr_number} + +states: + - name: analyze + type: agent + agent: analyzer + # Uses workflow git settings (sequential) + + - name: parallel_checks + type: parallel + git: + collaboration: parallel # Override for this section + branches: + - agent: security-scanner + - agent: perf-analyzer + transition: merge + + - name: merge + type: agent + agent: merge-bot + # Back to sequential (workflow default) +``` + +#### 5. Mental Model Summary + +``` +SESSION = Named, persistent workspace with git repo +├── Has a WORKING BRANCH (the main line of work) +├── Tracks which RUN currently holds the lock (sequential mode) +└── Persists until explicitly deleted + +SEQUENTIAL COLLABORATION +├── Agents work on same branch, same files +├── Each sees previous agent's commits +├── "Pick up where left off" semantics +├── Only ONE active run at a time (locked) +└── Default for workflows + +PARALLEL COLLABORATION +├── Each agent gets own branch from working branch +├── Multiple concurrent runs OK +├── Requires merge step to combine +└── Must be explicitly configured + +RUN = Single agent execution +├── In sequential: acquires session lock, works on working branch +├── In parallel: creates/uses own branch +└── On completion: commits changes, releases lock (if held) +``` + +### Implementation Plan + +| Phase | Description | Status | +|-------|-------------|--------| +| 2.1 | Add `ExecutionContext` to harness config | ✅ Done | +| 2.2 | Implement `ResolveWorkspacePath()` | ✅ Done | +| 2.3 | Add `git/cloner.go` for repo cloning | ✅ Done | +| 2.4 | Add session lock manager for sequential mode | ✅ Done | +| 2.5 | Wire into `executeWithAgenticHarness()` | ✅ Done | +| 2.6 | Add real-time streaming with full identifiers | ✅ Done | +| 2.7 | Session management commands (`stn session list/delete/cleanup/unlock/info`) | ✅ Done | +| 2.8 | Add CLI flags (`--session`, `--repo`, `--collaboration`) | Pending | +| 2.9 | Add workflow git/collaboration settings parsing | Pending | +| 2.10 | E2E tests for sequential continuation | Pending | +| 2.11 | E2E tests for parallel branches + merge | Pending | + +### Success Criteria (Phase 2) + +- [ ] **Sequential mode**: Agent B sees Agent A's commits on same branch +- [ ] **Parallel mode**: Each agent gets isolated branch, merge step combines +- [ ] **Session locking**: Sequential mode blocks concurrent runs on same session +- [ ] **Session persistence**: `--session` flag continues from previous run +- [ ] **Git repo cloning**: `--repo` flag clones repo into session workspace +- [ ] **Workflow integration**: `git.collaboration` setting propagates to steps +- [ ] **CLI commands**: `stn session list/delete/cleanup` work correctly +- [ ] **Lock recovery**: Stale locks (crashed runs) auto-expire after timeout +- [x] **Real-time streaming**: Events streamed with full ownership identifiers + +## Phase 2.6: Real-Time Streaming (COMPLETE) + +### Overview + +The streaming system enables real-time visibility into agent execution, similar to Claude Code/OpenCode. Events are published to NATS and can be consumed by Lighthouse for forwarding to CloudShip platform. + +### Architecture + +``` +AgenticExecutor.runLoop() + │ + ├─► run_start ──────────► StreamEvent + │ + ├─► For each step: + │ ├─► tool_start ─────► StreamEvent + │ ├─► tool_result ────► StreamEvent + │ └─► step_complete ──► StreamEvent + │ + └─► run_complete ───────► StreamEvent + │ + ▼ + NATS: station.{station_id}.run.{run_uuid}.stream + │ + ▼ + Lighthouse StreamConsumer (future) + │ + ▼ + CloudShip Platform → WebSocket → Browser UI +``` + +### Stream Event Structure + +Every event contains full ownership identifiers for correlation: + +```go +type Event struct { + // Ownership identifiers - who does this stream belong to? + StationRunID string // Local DB ID (e.g., "123") + RunUUID string // CloudShip's unique run identifier + WorkflowRunID string // Workflow run ID (when in workflow context) + SessionID string // Session ID for workspace persistence + AgentID string // Agent ID (e.g., "agent-1") + AgentName string // Human-readable agent name + StationID string // Station identifier + + // Event metadata + Seq int64 // Sequence number within the stream + Timestamp time.Time // Event timestamp + Type EventType // Event type + Data any // Event payload +} +``` + +### Event Types + +| Type | Data Structure | When Emitted | +|------|----------------|--------------| +| `run_start` | `RunStartData{AgentID, AgentName, Task, MaxSteps}` | Execution begins | +| `tool_start` | `ToolStartData{ToolName, ToolID, Input}` | Before tool execution | +| `tool_result` | `ToolResultData{ToolName, ToolID, Output, DurationMs, Error}` | After tool execution | +| `step_complete` | `StepCompleteData{StepNumber, TotalTokens, InputTokens, OutputTokens, FinishReason}` | After each LLM step | +| `run_complete` | `RunCompleteData{Success, TotalSteps, TotalTokens, DurationMs, FinishReason, Error}` | Execution ends | +| `token` | `TokenData{Content, Done}` | Streaming tokens (future) | +| `thinking` | `ThinkingData{Content}` | Extended thinking content (future) | +| `error` | `{error: string}` | Error occurred | + +### Correlation Matrix + +Different systems use different identifiers to correlate streams: + +| System | Uses | Purpose | +|--------|------|---------| +| CloudShip Platform | `run_uuid` | Correlate with their runs table | +| Station Local DB | `station_run_id` | Correlate with agent_runs table | +| Workflows | `workflow_run_id` | Group steps in same workflow | +| Sessions | `session_id` | Resume work in same workspace | +| Lighthouse | `station_id` + `run_uuid` | Route to correct station | +| NATS Subject | `station_id` + `run_uuid` | Message routing | + +### NATS Subject Pattern + +``` +station.{station_id}.run.{run_uuid}.stream +``` + +Subscription patterns: +- All streams from a station: `station.my-station.run.*.stream` +- Specific run: `station.my-station.run.abc-123.stream` +- All stations: `station.*.run.*.stream` + +### Configuration + +```yaml +harness: + streaming: + enabled: true +``` + +### Publisher Implementations + +| Publisher | Use Case | +|-----------|----------| +| `ChannelPublisher` | Local/testing - publishes to Go channel | +| `NATSPublisher` | Production - publishes to NATS (with optional JetStream) | +| `MultiPublisher` | Combine multiple publishers | +| `NoOpPublisher` | Disabled streaming | + +### Component Structure + +``` +pkg/harness/stream/ +├── publisher.go # Event types, StreamContext, Publisher interface +└── nats_publisher.go # NATS/JetStream publisher implementation +``` + +### Usage in Executor + +```go +// ExecuteOptions includes all stream identifiers +execOpts := harness.ExecuteOptions{ + StationRunID: fmt.Sprintf("%d", agent.ID), + RunUUID: runUUID, + WorkflowRunID: workflowRunID, + SessionID: workflowRunID, + AgentName: agent.Name, + StationID: stationID, +} + +// Publisher is configured via option +executor := harness.NewAgenticExecutor( + genkitApp, + harnessConfig, + agentHarnessConfig, + harness.WithStreamPublisher(natsPublisher), +) + +// Execute with streaming +result, err := executor.ExecuteWithOptions(ctx, agentID, task, tools, execOpts) +``` + +### Example Event Sequence + +```json +{"station_run_id":"123","run_uuid":"abc-def","workflow_run_id":"wf-001","agent_id":"5","agent_name":"code-analyzer","station_id":"stn-xyz","seq":1,"type":"run_start","data":{"agent_id":"5","agent_name":"code-analyzer","task":"Analyze main.go","max_steps":50}} + +{"station_run_id":"123","run_uuid":"abc-def","workflow_run_id":"wf-001","agent_id":"5","agent_name":"code-analyzer","station_id":"stn-xyz","seq":2,"type":"tool_start","data":{"tool_name":"read","tool_id":"read-123","input":{"path":"main.go"}}} + +{"station_run_id":"123","run_uuid":"abc-def","workflow_run_id":"wf-001","agent_id":"5","agent_name":"code-analyzer","station_id":"stn-xyz","seq":3,"type":"tool_result","data":{"tool_name":"read","tool_id":"read-123","output":"package main...","duration_ms":5}} + +{"station_run_id":"123","run_uuid":"abc-def","workflow_run_id":"wf-001","agent_id":"5","agent_name":"code-analyzer","station_id":"stn-xyz","seq":4,"type":"step_complete","data":{"step_number":1,"total_tokens":1500,"input_tokens":1000,"output_tokens":500,"finish_reason":"tool_use"}} + +{"station_run_id":"123","run_uuid":"abc-def","workflow_run_id":"wf-001","agent_id":"5","agent_name":"code-analyzer","station_id":"stn-xyz","seq":5,"type":"run_complete","data":{"success":true,"total_steps":3,"total_tokens":4500,"duration_ms":12000,"finish_reason":"agent_done"}} +``` + +### Future Work + +- [ ] Lighthouse StreamConsumer to forward events to CloudShip +- [ ] Token-by-token streaming for real-time UI updates +- [ ] Extended thinking content streaming +- [ ] WebSocket adapter for direct browser connections +- [ ] Stream replay from NATS JetStream + +## Phase 3: Sandbox Isolation Strategies (COMPLETE) + +### Problem Statement + +The current harness executes code directly on the host system. This presents security risks: + +1. **Untrusted code execution** - Agent-generated bash commands run with user privileges +2. **File system access** - Agents can read/write anywhere the user has access +3. **Network access** - Agents can make arbitrary network requests +4. **Resource exhaustion** - No limits on CPU, memory, or disk usage +5. **Persistence** - Changes persist on host even after errors + +### Sandbox Strategies + +| Strategy | Isolation Level | Startup Time | Use Case | Status | +|----------|-----------------|--------------|----------|--------| +| **Host** | None | Instant | Development, trusted agents | ✅ Complete | +| **Docker** | Container | ~1-2s | Production, untrusted code | ✅ Complete | +| **E2B** | Cloud sandbox | ~2-3s | Serverless, managed infra | ✅ Complete | +| **Firecracker** | microVM | ~200ms | High security, multi-tenant | Planned | +| **gVisor** | Kernel sandbox | ~500ms | Balance of security/performance | Planned | +| **WASM** | Process sandbox | ~50ms | Lightweight, fast iteration | Planned | + +### Proposed Architecture + +``` +┌─────────────────────────────────────────────────────────────────┐ +│ SANDBOX CONFIGURATION │ +│ (per-agent or global config) │ +├─────────────────────────────────────────────────────────────────┤ +│ sandbox: │ +│ mode: docker | firecracker | gvisor | wasm | host │ +│ image: station-sandbox:latest │ +│ resources: │ +│ cpu: 2 │ +│ memory: 4Gi │ +│ disk: 10Gi │ +│ network: │ +│ enabled: false │ +│ allowed_hosts: [github.com, api.openai.com] │ +│ filesystem: │ +│ read_only: [/etc, /usr] │ +│ read_write: [/workspace] │ +│ denied: [/etc/passwd, ~/.ssh] │ +│ timeout: 30m │ +└─────────────────────────────────────────────────────────────────┘ + │ + ▼ +┌─────────────────────────────────────────────────────────────────┐ +│ SANDBOX MANAGER │ +│ │ +│ Interface: │ +│ - Create(config) → SandboxInstance │ +│ - Execute(instance, command) → Result │ +│ - ReadFile(instance, path) → Content │ +│ - WriteFile(instance, path, content) → Error │ +│ - Destroy(instance) → Error │ +│ │ +│ Implementations: │ +│ - HostSandbox (passthrough, no isolation) │ +│ - DockerSandbox (container-based) │ +│ - FirecrackerSandbox (microVM-based) │ +│ - WASMSandbox (WebAssembly-based) │ +└─────────────────────────────────────────────────────────────────┘ +``` + +### Implementation Plan + +| Phase | Description | Priority | Status | +|-------|-------------|----------|--------| +| 3.1 | Define `Sandbox` interface and `SandboxConfig` | High | ✅ Complete | +| 3.2 | Implement `HostSandbox` (current behavior, passthrough) | High | ✅ Complete | +| 3.3 | Implement `DockerSandbox` with resource limits | High | ✅ Complete | +| 3.4 | Implement `E2BSandbox` with cloud execution | High | ✅ Complete | +| 3.5 | Add `StreamingSandbox` interface with real-time output | High | ✅ Complete | +| 3.6 | Implement `RuntimeResolver` for auto-detection | High | ✅ Complete | +| 3.7 | Wire sandbox into `AgenticExecutor` lifecycle | High | ✅ Complete | +| 3.8 | Add network filtering (allowed hosts list) | Medium | Pending | +| 3.9 | Add filesystem ACLs (read-only, read-write, denied paths) | Medium | Pending | +| 3.10 | Implement `FirecrackerSandbox` for high-security use cases | Low | Pending | +| 3.11 | Implement `WASMSandbox` for lightweight isolation | Low | Pending | +| 3.12 | Add sandbox metrics (CPU, memory, I/O usage) | Medium | Pending | +| 3.13 | Add sandbox cleanup/garbage collection | Medium | Pending | +| 3.14 | E2E tests for each sandbox mode | High | ✅ Complete (Host, Docker, E2B) | + +### Docker Sandbox Design + +```go +type DockerSandbox struct { + containerID string + config SandboxConfig + client *docker.Client +} + +// Creates container with: +// - Read-only root filesystem +// - Mounted workspace volume +// - Resource limits (CPU, memory) +// - No network (unless explicitly allowed) +// - Dropped capabilities +func (s *DockerSandbox) Create(ctx context.Context) error + +// Executes command via docker exec +func (s *DockerSandbox) Execute(ctx context.Context, cmd string) (string, error) + +// Copies file from container +func (s *DockerSandbox) ReadFile(ctx context.Context, path string) ([]byte, error) + +// Copies file into container +func (s *DockerSandbox) WriteFile(ctx context.Context, path string, content []byte) error +``` + +### Configuration Examples + +**Development (no isolation):** +```yaml +harness: + sandbox: + mode: host +``` + +**Production (Docker isolation):** +```yaml +harness: + sandbox: + mode: docker + image: station-sandbox:latest + resources: + cpu: 2 + memory: 4Gi + network: + enabled: false + timeout: 30m +``` + +**High-security (Firecracker microVM):** +```yaml +harness: + sandbox: + mode: firecracker + kernel: /var/lib/firecracker/vmlinux + rootfs: /var/lib/firecracker/rootfs.ext4 + resources: + vcpu: 2 + memory: 2048 + network: + enabled: true + allowed_hosts: + - "*.github.com" + - "api.openai.com" +``` + +### Success Criteria (Phase 3) + +- [x] **Host mode**: Current behavior preserved, zero overhead +- [x] **Docker mode**: Agents execute in isolated containers +- [x] **E2B mode**: Agents execute in cloud sandboxes with streaming +- [x] **Streaming support**: Real-time stdout/stderr via `StreamingSandbox` interface +- [x] **RuntimeResolver**: Auto-detects host capabilities, falls back to configured isolation +- [x] **Executor integration**: Sandbox initialized in setup(), destroyed in cleanup() +- [x] **Resource limits**: CPU/memory/disk limits enforced (Docker, E2B) +- [ ] **Network filtering**: Only allowed hosts reachable (future) +- [ ] **Filesystem ACLs**: Denied paths inaccessible (future) +- [x] **Cleanup**: Sandbox destroyed on completion/timeout +- [ ] **Metrics**: Resource usage tracked per execution (future) +- [x] **E2E tests**: Host, Docker, E2B sandbox modes tested (25 tests pass) + +### Component Structure (Phase 3) + +``` +pkg/harness/sandbox/ +├── sandbox.go # Sandbox interface, Config, Factory, RuntimeResolver +├── host.go # HostSandbox - passthrough execution +├── docker.go # DockerSandbox - container isolation +├── e2b.go # E2BSandbox - cloud sandbox with streaming +├── sandbox_test.go # Unit tests (13 tests) +├── e2b_test.go # E2B unit tests (12 tests) +└── e2b/ + ├── client.go # REST client for E2B API + ├── streaming.go # gRPC/Connect streaming client + └── proto/ + ├── process.proto # E2B envd process service definition + └── gen/ # Generated protobuf and Connect code +``` + +### Key Interfaces + +```go +// Sandbox provides isolated execution environment +type Sandbox interface { + ID() string + Exec(ctx context.Context, command string, args ...string) (string, error) + ReadFile(ctx context.Context, path string) ([]byte, error) + WriteFile(ctx context.Context, path string, content []byte) error + Destroy(ctx context.Context) error +} + +// StreamingSandbox adds real-time output streaming +type StreamingSandbox interface { + Sandbox + ExecStream(ctx context.Context, command string, args []string, opts ExecOptions) (ProcessHandle, error) + ListProcesses(ctx context.Context) ([]ProcessInfo, error) + KillProcess(ctx context.Context, pid int) error +} + +// RuntimeResolver auto-detects sandbox mode based on host capabilities +type RuntimeResolver struct { + FallbackBackend IsolationBackend // docker | e2b +} + +func (r *RuntimeResolver) Resolve(cfg Config, reqs RuntimeRequirements) (Config, error) +``` + +### Runtime Resolution Flow + +``` +┌─────────────────────────────────────────────────────────────┐ +│ User Configuration │ +│ isolation_backend: docker | e2b ← User's explicit choice │ +└─────────────────────────────────────────────────────────────┘ + │ + ▼ +┌─────────────────────────────────────────────────────────────┐ +│ RuntimeResolver │ +│ if cfg.Mode == ModeAuto: │ +│ if host has required commands → ModeHost │ +│ else → use configured isolation_backend │ +└─────────────────────────────────────────────────────────────┘ + │ + ▼ +┌─────────────────────────────────────────────────────────────┐ +│ Factory.Create(cfg) │ +│ ModeHost → HostSandbox (passthrough) │ +│ ModeDocker → DockerSandbox (container) │ +│ ModeE2B → E2BSandbox (cloud) │ +└─────────────────────────────────────────────────────────────┘ +``` + +**Key principle:** Same agent config works everywhere. Station adapts based on environment. + +## Phase 4: Session Persistence & REPL Support (COMPLETE) + +### Problem Statement + +The harness executes agents statelessly - each execution starts fresh with no memory of previous interactions. For REPL-style workflows, users need: + +1. **Message history persistence** - Continue conversations across multiple inputs +2. **Sandbox persistence** - Keep the same files/environment between interactions +3. **Session lifecycle management** - Start, pause, resume, and end sessions + +### Architecture + +``` +┌─────────────────────────────────────────────────────────────────┐ +│ REPL Session Flow: │ +├─────────────────────────────────────────────────────────────────┤ +│ 1. StartSession(sessionID) │ +│ ├── GetOrCreate session workspace │ +│ ├── AcquireLock (prevent concurrent access) │ +│ ├── Initialize workspace directory │ +│ └── Create sandbox (keep alive for session duration) │ +├─────────────────────────────────────────────────────────────────┤ +│ 2. Execute(task) - repeat for each user input │ +│ ├── Load history from .history.json │ +│ ├── Create AgenticExecutor with loaded history │ +│ ├── Run task │ +│ ├── Append new messages to history │ +│ └── Save history to disk │ +├─────────────────────────────────────────────────────────────────┤ +│ 3. EndSession() │ +│ ├── Destroy sandbox (or keep for future resume) │ +│ └── ReleaseLock │ +└─────────────────────────────────────────────────────────────────┘ +``` + +### Session Storage + +Files stored per session: +``` +~/.config/station/workspace/session/{session-id}/ +├── .session.meta # Session metadata (created_at, total_runs) +├── .session.lock # Lock file (PID, expires_at) +├── .history.json # Message history (all user/assistant messages) +└── (workspace files created by agent) +``` + +### Key Components + +#### HistoryStore (`pkg/harness/session/history_store.go`) + +Persists message history for sessions enabling REPL-style conversations: + +```go +type HistoryStore struct { + basePath string + mu sync.RWMutex +} + +type StoredMessage struct { + Role string `json:"role"` + Content string `json:"content"` + ToolCalls []StoredToolCall `json:"tool_calls,omitempty"` + Metadata map[string]interface{} `json:"metadata,omitempty"` + Timestamp time.Time `json:"timestamp"` +} + +type SessionHistory struct { + SessionID string `json:"session_id"` + Messages []StoredMessage `json:"messages"` + TotalTokens int `json:"total_tokens"` + CreatedAt time.Time `json:"created_at"` + UpdatedAt time.Time `json:"updated_at"` +} + +// Load loads message history for a session +func (h *HistoryStore) Load(sessionID string) (*SessionHistory, error) + +// Save persists message history for a session +func (h *HistoryStore) Save(sessionID string, history *SessionHistory) error + +// Append adds messages to a session's history +func (h *HistoryStore) Append(sessionID string, messages []StoredMessage) error + +// ToAIMessages converts stored messages to ai.Message format +func (h *SessionHistory) ToAIMessages() []*ai.Message +``` + +#### SessionExecutor (`pkg/harness/session/session_executor.go`) + +Wraps AgenticExecutor with session persistence: + +```go +type SessionExecutor struct { + genkitApp *genkit.Genkit + sessionManager *Manager + historyStore *HistoryStore + // ... config fields + + // Active session state + currentSessionID string + currentSession *Session + currentWorkspace harness.WorkspaceManager + currentSandbox sandbox.Sandbox +} + +// StartSession starts or resumes a session +func (se *SessionExecutor) StartSession(ctx context.Context, sessionID string, repoURL string, branch string) error + +// Execute runs a task within the current session, maintaining history +func (se *SessionExecutor) Execute(ctx context.Context, task string) (*harness.ExecutionResult, error) + +// EndSession ends the current session, cleaning up resources +func (se *SessionExecutor) EndSession(ctx context.Context) error +``` + +#### AgenticExecutor History Support + +The AgenticExecutor now supports initial history injection: + +```go +// WithInitialHistory sets the initial message history for session persistence +func WithInitialHistory(history []*ai.Message) ExecutorOption + +// ExecutionResult now includes history for persistence +type ExecutionResult struct { + // ... existing fields + History []*ai.Message `json:"-"` // Full message history for session persistence +} +``` + +### Usage Example + +```go +// Create session executor +executor := session.NewSessionExecutor(genkitApp, session.SessionExecutorConfig{ + BasePath: "~/.config/station/workspace", + ModelName: "openai/gpt-4o-mini", + SystemPrompt: "You are a helpful coding assistant.", +}) + +// Start session +if err := executor.StartSession(ctx, "my-coding-session", "", ""); err != nil { + log.Fatal(err) +} +defer executor.EndSession(ctx) + +// First interaction +result1, _ := executor.Execute(ctx, "Create a hello.py file that prints 'Hello World'") + +// Second interaction - agent has full context of previous work +result2, _ := executor.Execute(ctx, "Now add a function that takes a name parameter") + +// Third interaction - continues with full history +result3, _ := executor.Execute(ctx, "Add tests for the function") +``` + +### CLI Integration (Future) + +```bash +# Start REPL session +stn repl --session "my-project" --agent "coder" + +# Inside REPL: +> Create a REST API with user CRUD operations +# Agent creates files, you see the output +> Add authentication middleware +# Agent continues, aware of previous work +> exit + +# Later, resume the session +stn repl --session "my-project" --agent "coder" +# Agent has full context of all previous work +``` + +### Implementation Status + +| Phase | Description | Status | +|-------|-------------|--------| +| 4.1 | Define `StoredMessage` and `SessionHistory` types | ✅ Complete | +| 4.2 | Implement `HistoryStore` for JSON persistence | ✅ Complete | +| 4.3 | Add `WithInitialHistory` option to `AgenticExecutor` | ✅ Complete | +| 4.4 | Return `History` in `ExecutionResult` | ✅ Complete | +| 4.5 | Implement `SessionExecutor` wrapper | ✅ Complete | +| 4.6 | Add session persistence tests | ✅ Complete | +| 4.7 | Add CLI REPL command (`stn repl`) | Pending | +| 4.8 | Add session list/delete CLI commands | Pending | + +### Success Criteria (Phase 4) + +- [x] **HistoryStore**: Can save/load/append message history to disk +- [x] **SessionExecutor**: Wraps AgenticExecutor with session lifecycle +- [x] **Initial history injection**: `WithInitialHistory` option works +- [x] **History in result**: ExecutionResult includes full history +- [x] **Unit tests**: All session persistence tests pass +- [ ] **CLI REPL**: `stn repl` command for interactive sessions (future) +- [ ] **E2E test**: Multi-turn conversation with real LLM (future) + +### Component Structure (Phase 4) + +``` +pkg/harness/session/ +├── manager.go # Session lifecycle management +├── history_store.go # Message history persistence (NEW) +├── session_executor.go # Session-aware executor wrapper (NEW) +└── session_executor_test.go # Session persistence tests (NEW) +``` + +### Test Results (2025-01-26) + +All session tests pass: + +| Test | Description | Status | +|------|-------------|--------| +| `TestHistoryStore_SaveLoad` | Save and load history | ✅ PASS | +| `TestHistoryStore_Append` | Append messages across calls | ✅ PASS | +| `TestHistoryStore_Clear` | Clear history | ✅ PASS | +| `TestSessionManager_Integration` | Full session lifecycle | ✅ PASS | + +Run tests: +```bash +go test ./pkg/harness/session/... -v -count=1 +``` + +## References + +- `internal/services/agent_execution_engine.go` - Integration point +- `pkg/harness/` - Harness implementation +- `internal/workflows/runtime/` - Workflow execution system +- `internal/lattice/work/` - Existing NATS/state infrastructure +- `internal/services/sandbox_session_manager.go` - Existing session pattern diff --git a/docs/testing/HARNESS_TESTING_OPS.md b/docs/testing/HARNESS_TESTING_OPS.md new file mode 100644 index 00000000..8d957cde --- /dev/null +++ b/docs/testing/HARNESS_TESTING_OPS.md @@ -0,0 +1,377 @@ +# Agentic Harness Testing Operations Guide + +**Last Verified**: 2025-01-26 +**Test Results**: 90+ tests passing across harness packages (including Docker sandbox fixes) + +## Quick Reference + +```bash +# Run all harness unit tests +go test ./pkg/harness/... -v -count=1 + +# Run with race detection +go test ./pkg/harness/... -race -v + +# Run E2E tests (requires real LLM) +HARNESS_E2E_TEST=1 go test ./pkg/harness/... -v -run E2E +``` + +--- + +## Test Coverage Summary + +| Package | Tests | Status | Notes | +|---------|-------|--------|-------| +| `pkg/harness` | 23 | ✅ PASS | Core executor, config, doom loop | +| `pkg/harness/prompt` | 20 | ✅ PASS | Agent config parsing, frontmatter | +| `pkg/harness/skills` | 9 | ✅ PASS | Skills middleware, parsing | +| `pkg/harness/memory` | 6 | ✅ PASS | Memory config, middleware | +| `pkg/harness/hooks` | 8 | ✅ PASS | Permission checking, wildcard matching | +| `pkg/harness/stream` | 10 | ✅ PASS | Event streaming, publishers | +| `pkg/harness/workspace` | 10 | ✅ PASS | Host workspace operations | +| `pkg/harness/git` | 5 | ✅ PASS | Git manager integration | +| `pkg/harness/nats` | 4 | ✅ PASS | NATS adapters, lattice | +| `pkg/harness/sandbox` | 25+ | ✅ PASS | Docker sandbox fixed (uses docker cp for file ops) | + +--- + +## Unit Tests + +### Core Harness Tests (`pkg/harness/`) + +```bash +go test ./pkg/harness/ -v -count=1 +``` + +**What's tested:** +- `TestCompactor_*` - Context compaction, token counting, history truncation +- `TestDoomLoopDetector_*` - Doom loop detection via SHA256 hashing +- `TestAgenticExecutor_*` - Executor initialization, sandbox options +- `TestDefaultHarnessConfig` - Default config values +- `TestHookResult_*` - Pre/post hook execution priority + +### Prompt/Config Tests (`pkg/harness/prompt/`) + +```bash +go test ./pkg/harness/prompt/... -v -count=1 +``` + +**What's tested:** +- `TestParseAgentConfig_*` - Frontmatter parsing (model, tools, harness, sandbox) +- `TestAgentConfig_Defaults` - Default value injection +- `TestBuilder_*` - System prompt building +- `TestAgentConfig_GetSkillSources_*` - Skill source resolution + +### Permission/Hook Tests + +```bash +go test ./pkg/harness/ -v -run "Hook\|Match\|Permission" +``` + +**What's tested:** +- `TestMatchWildcard` - Wildcard pattern matching (`git *`, `rm -rf *`) +- `TestCheckBashPermission` - Allow/deny/ask rule evaluation +- `TestHookRegistry_*` - Hook priority (Interrupt > Block > Continue) + +--- + +## E2E Tests + +E2E tests require a real LLM API key and are skipped by default. + +### Enable E2E Tests + +```bash +# Set required environment variables +export HARNESS_E2E_TEST=1 +export OPENAI_API_KEY="sk-..." # Or other provider + +# Run all E2E tests +go test ./pkg/harness/... -v -run E2E -timeout 10m +``` + +### Available E2E Tests + +| Test | What It Validates | Requirements | +|------|-------------------|--------------| +| `TestAgenticExecutor_E2E_RealLLM` | Basic file creation task | LLM API | +| `TestAgenticExecutor_E2E_MultiStep` | Multi-step task execution | LLM API | +| `TestAgenticExecutor_E2E_Compaction` | Context compaction triggers | LLM API | +| `TestAgenticExecutor_E2E_WorkflowSimulation` | Multi-agent handoff | LLM API | +| `TestAgenticExecutor_E2E_DockerSandbox` | Docker sandbox isolation | LLM API + Docker | +| `TestAgenticExecutor_E2E_DockerWorkflowHandoff` | Workspace persistence across containers | LLM API + Docker | + +### Docker E2E Tests + +```bash +# Requires Docker daemon running +HARNESS_E2E_TEST=1 go test ./pkg/harness/... -v -run Docker -timeout 10m +``` + +--- + +## Sandbox Tests + +### Host Sandbox (Always Works) + +```bash +go test ./pkg/harness/sandbox/... -v -run Host +``` + +### Docker Sandbox (Requires Docker) + +```bash +# Run Docker-specific tests +go test ./pkg/harness/sandbox/... -v -run Docker + +# Known failures in restricted environments: +# - TestDockerSandbox_CopyOut_AutoCreatesContainer +# - TestDockerSandbox_FileOperations +# These require Docker with proper permissions for container file ops +``` + +### E2B Sandbox (Mocked by Default) + +```bash +# E2B tests use mocks - no API key needed for unit tests +go test ./pkg/harness/sandbox/... -v -run E2B + +# Real E2B integration (requires API key) +E2B_API_KEY="..." go test ./pkg/harness/sandbox/... -v -run E2B -tags=integration +``` + +--- + +## CI/CD Integration + +### GitHub Actions Workflow + +```yaml +name: Harness Tests + +on: + push: + paths: + - 'pkg/harness/**' + pull_request: + paths: + - 'pkg/harness/**' + +jobs: + unit-tests: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-go@v5 + with: + go-version: '1.22' + + - name: Run Unit Tests + run: go test ./pkg/harness/... -v -count=1 -race + + - name: Run Unit Tests (exclude Docker) + run: go test ./pkg/harness/... -v -count=1 -skip "Docker" + + e2e-tests: + runs-on: ubuntu-latest + if: github.event_name == 'push' && github.ref == 'refs/heads/main' + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-go@v5 + with: + go-version: '1.22' + + - name: Run E2E Tests + env: + HARNESS_E2E_TEST: "1" + OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} + run: go test ./pkg/harness/... -v -run E2E -timeout 15m +``` + +--- + +## Environment Variables + +| Variable | Purpose | Required For | +|----------|---------|--------------| +| `HARNESS_E2E_TEST=1` | Enable E2E tests | All E2E tests | +| `OPENAI_API_KEY` | OpenAI API access | E2E with GPT models | +| `ANTHROPIC_API_KEY` | Anthropic API access | E2E with Claude models | +| `GEMINI_API_KEY` | Google AI access | E2E with Gemini models | +| `E2B_API_KEY` | E2B cloud sandbox | E2B integration tests | + +--- + +## Debugging Test Failures + +### Common Issues + +**1. Docker tests fail with permission errors** +``` +WriteFile failed: failed to write file /workspace/test.txt: exit status 1 +``` +- **Cause**: Docker container can't write to mounted volume +- **Fix**: Check Docker socket permissions, try `docker run --privileged` + +**2. E2E tests timeout** +``` +panic: test timed out after 2m0s +``` +- **Cause**: LLM API slow or rate limited +- **Fix**: Increase timeout: `go test -timeout 10m` + +**3. Doom loop detection triggers incorrectly** +``` +Doom loop detected: tool 'bash' has been called 3 times +``` +- **Cause**: Test task causes repeated tool calls +- **Fix**: Increase threshold in test config: `DoomLoopThreshold: 5` + +### Verbose Debugging + +```bash +# Enable all logging +go test ./pkg/harness/... -v -count=1 2>&1 | tee test-output.log + +# Run single test with verbose output +go test ./pkg/harness/ -v -run TestDoomLoopDetector_DetectsLoop -count=1 +``` + +--- + +## Test Data Locations + +| Data | Location | +|------|----------| +| E2E prompts | `/tmp/harness-e2e-prompts/` | +| Temp workspaces | `/tmp/harness-e2e-*/` | +| Test configs | Inline in test files | + +--- + +## Session Persistence (NEW) + +Session persistence enables REPL-style prolonged conversations with message history. + +### Session Tests + +```bash +go test ./pkg/harness/session/... -v -count=1 +``` + +| Test | What It Validates | +|------|-------------------| +| `TestHistoryStore_SaveLoad` | Message history save/load to disk | +| `TestHistoryStore_Append` | Appending messages across calls | +| `TestHistoryStore_Clear` | Clearing history | +| `TestSessionManager_Integration` | Full session lifecycle with history persistence | + +### How Session Persistence Works + +``` +REPL Session Flow: +┌─────────────────────────────────────────────────────────────────┐ +│ 1. StartSession(sessionID) │ +│ ├── GetOrCreate session workspace │ +│ ├── AcquireLock (prevent concurrent access) │ +│ ├── Initialize workspace directory │ +│ └── Create sandbox (keep alive for session duration) │ +├─────────────────────────────────────────────────────────────────┤ +│ 2. Execute(task) - repeat for each user input │ +│ ├── Load history from .history.json │ +│ ├── Create AgenticExecutor with loaded history │ +│ ├── Run task │ +│ ├── Append new messages to history │ +│ └── Save history to disk │ +├─────────────────────────────────────────────────────────────────┤ +│ 3. EndSession() │ +│ ├── Destroy sandbox (or keep for future resume) │ +│ └── ReleaseLock │ +└─────────────────────────────────────────────────────────────────┘ + +Files stored per session: +~/.config/station/workspace/session/{session-id}/ +├── .session.meta # Session metadata (created_at, total_runs) +├── .session.lock # Lock file (PID, expires_at) +├── .history.json # Message history (all user/assistant messages) +└── (workspace files created by agent) +``` + +--- + +## Coverage Gaps (TODO) + +The following areas need additional test coverage: + +| Area | Current State | Priority | +|------|---------------|----------| +| Session → AgenticExecutor history injection | ✅ Complete | ~~HIGH~~ | +| REPL Session Persistence | ✅ Complete | ~~HIGH~~ | +| Workflow approval integration | No tests | HIGH | +| NATS streaming E2E | Mocked only | MEDIUM | +| Lattice multi-station handoff | No tests | MEDIUM | +| Memory injection E2E | No tests | MEDIUM | +| Skills progressive disclosure | Basic tests | LOW | + +### Recommended New Tests + +1. **Workflow Approval E2E** + ```go + func TestWorkflowApproval_HumanInTheLoop(t *testing.T) { + // Test that HookInterrupt properly pauses execution + // Test approval/rejection flow + } + ``` + +2. **NATS Streaming E2E** + ```go + func TestNATSStreaming_EventDelivery(t *testing.T) { + // Test real NATS event publishing + // Test subscriber receives events + } + ``` + +3. **Memory Context Injection** + ```go + func TestMemoryInjection_E2E(t *testing.T) { + // Test AGENTS.md loading + // Test CloudShip memory API integration + } + ``` + +--- + +## Running Tests Locally + +### Full Test Suite (Fast) + +```bash +# Excludes E2E and Docker tests +go test ./pkg/harness/... -v -count=1 -skip "E2E|Docker" +``` + +### Full Test Suite (Complete) + +```bash +# All tests including Docker (requires Docker) +HARNESS_E2E_TEST=1 OPENAI_API_KEY="sk-..." \ + go test ./pkg/harness/... -v -timeout 15m +``` + +### Quick Smoke Test + +```bash +# Just the core logic +go test ./pkg/harness/ -v -count=1 +``` + +--- + +## Maintenance Notes + +- **Update this doc** when adding new test files +- **Tag flaky tests** with `t.Skip()` and reason +- **Docker tests** may fail in CI without privileged mode +- **E2E tests** cost money (LLM API calls) - run sparingly diff --git a/environments/default/agents/harness-test.prompt b/environments/default/agents/harness-test.prompt new file mode 100644 index 00000000..efeda1eb --- /dev/null +++ b/environments/default/agents/harness-test.prompt @@ -0,0 +1,27 @@ +--- +name: harness-test +description: Test agent for agentic harness integration +harness: agentic +harness_config: + max_steps: 25 + doom_loop_threshold: 5 + timeout: 15m + sandbox: + mode: host +tools: + - read + - write + - bash + - glob + - grep +--- +You are a test agent running with the agentic harness. + +Your task is to help with simple file operations and shell commands in the workspace. + +When given a task: +1. Understand what is being asked +2. Use the appropriate tools to complete the task +3. Report back with a clear summary + +Be concise and efficient. diff --git a/environments/default/agents/incident-investigator.prompt b/environments/default/agents/incident-investigator.prompt new file mode 100644 index 00000000..0c2da6ca --- /dev/null +++ b/environments/default/agents/incident-investigator.prompt @@ -0,0 +1,98 @@ +--- +model: openai/gpt-4o +metadata: + name: incident-investigator + description: Expert incident investigator for cloud-native systems + tags: + - sre + - incidents + - kubernetes +tools: + - bash + - read_file + - write_file + - edit_file + - glob + - grep + +harness: + enabled: true + max_steps: 100 + timeout: 60m + doom_loop: + threshold: 3 + action: summarize + compaction: + enabled: true + threshold: 0.85 + protect_tokens: 40000 + history_offload: true + progress: + file: progress.md + sync_interval: 5s + +workspace: + mode: sandbox + scope: session + cleanup: on_success + +sandbox: + mode: docker + image: python:3.12-slim + timeout: 30m + resources: + cpu: "2" + memory: 4g + network: + enabled: true + allowed_hosts: + - "*.github.com" + - "pypi.org" + +skills: + sources: + - /skills/user/ + - /skills/project/ + +memory: + sources: + - ~/.config/station/AGENTS.md + - .station/AGENTS.md +--- + +{{role "system"}} +You are an expert incident investigator for cloud-native systems. Your task is to analyze incidents, identify root causes, and recommend remediation. + +## Your Capabilities + +You have access to tools for: +- **Kubernetes inspection**: Use bash with kubectl to examine pods, services, deployments, events +- **Metrics analysis**: Query time-series data and logs +- **Log analysis**: Use bash with grep, jq for log parsing +- **File operations**: Read, write, and edit files in your workspace + +## Investigation Methodology + +1. **Gather Context**: Check pod status, recent events, service health +2. **Analyze Metrics**: Look for anomalies in CPU, memory, error rates +3. **Trace Timeline**: Build chronological sequence of events +4. **Identify Root Cause**: Correlate findings to determine cause +5. **Document Findings**: Write structured output with recommendations + +## Skills System + +You have access to specialized skills that provide domain-specific workflows. Check the skills section in your system prompt - if a skill matches your task, read its full instructions before proceeding. + +## Memory Context + +Pay attention to any project-specific context loaded from AGENTS.md files. This contains important information about the environment, conventions, and user preferences. + +## Progress Tracking + +Update `progress.md` in your workspace as you work: +- Mark completed investigation steps +- Note key findings as you discover them +- Track what's left to investigate + +{{role "user"}} +{{userInput}} diff --git a/environments/default/memory/AGENTS.md b/environments/default/memory/AGENTS.md new file mode 100644 index 00000000..2c61220e --- /dev/null +++ b/environments/default/memory/AGENTS.md @@ -0,0 +1,29 @@ +# Project Context + +## Overview +This is a DevOps automation platform built with Go and React. + +## Build & Test Commands +- `make build` - Build all components +- `make test` - Run unit tests +- `make e2e` - Run E2E tests with real LLM + +## Architecture Notes +- Backend: Go with Genkit for LLM orchestration +- Frontend: React with TypeScript +- Database: PostgreSQL for runs, NATS KV for sessions + +## Code Style +- Use `slog` for structured logging +- Error wrapping with `fmt.Errorf("context: %w", err)` +- Table-driven tests preferred + +## User Preferences +- Prefers concise explanations +- Uses dark mode +- Timezone: PST + +## Learned Patterns +- When querying Prometheus, always use the `prometheus-query` MCP tool +- For Kubernetes issues, check pod events first via `kubectl describe` +- User prefers markdown tables over bullet lists for comparison data diff --git a/environments/default/skills/user/code-review/SKILL.md b/environments/default/skills/user/code-review/SKILL.md new file mode 100644 index 00000000..35104ef2 --- /dev/null +++ b/environments/default/skills/user/code-review/SKILL.md @@ -0,0 +1,74 @@ +--- +name: code-review +description: Systematic code review process with focus on quality, security, and best practices +allowed-tools: + - read_file + - grep + - glob + - bash +triggers: + - review + - code quality + - pull request +--- + +# Code Review Skill + +## When to Use +- User asks for code review or feedback +- Reviewing a pull request or diff +- Analyzing code quality or security + +## Review Checklist + +### 1. Functionality +- [ ] Does the code do what it's supposed to? +- [ ] Are edge cases handled? +- [ ] Is error handling appropriate? + +### 2. Code Quality +- [ ] Is the code readable and maintainable? +- [ ] Are functions and variables named clearly? +- [ ] Is there unnecessary duplication? +- [ ] Are comments helpful and accurate? + +### 3. Security +- [ ] Input validation present? +- [ ] No hardcoded secrets? +- [ ] SQL injection prevention? +- [ ] XSS prevention? + +### 4. Performance +- [ ] Efficient algorithms used? +- [ ] No N+1 query problems? +- [ ] Resources properly managed? + +### 5. Testing +- [ ] Tests cover main functionality? +- [ ] Edge cases tested? +- [ ] Tests are readable? + +## Output Format + +For each finding, provide: +1. **Location**: File and line number +2. **Severity**: Critical / High / Medium / Low +3. **Issue**: Clear description of the problem +4. **Suggestion**: How to fix it +5. **Example**: Code snippet showing the fix (if applicable) + +## Example Review Comment + +**Location**: `src/api/users.go:42` +**Severity**: High +**Issue**: SQL query built using string concatenation +**Suggestion**: Use parameterized queries to prevent SQL injection +**Example**: +```go +// Before (vulnerable) +query := "SELECT * FROM users WHERE id = " + userID + +// After (safe) +query := "SELECT * FROM users WHERE id = ?" +db.Query(query, userID) +``` diff --git a/environments/default/skills/user/web-research/SKILL.md b/environments/default/skills/user/web-research/SKILL.md new file mode 100644 index 00000000..2560a2bf --- /dev/null +++ b/environments/default/skills/user/web-research/SKILL.md @@ -0,0 +1,67 @@ +--- +name: web-research +description: Structured approach to conducting thorough web research with source citation +license: MIT +allowed-tools: + - web_search + - web_fetch + - read_file + - write_file +triggers: + - research + - search + - find information +--- + +# Web Research Skill + +## When to Use +- User asks you to research a topic +- You need to gather information from multiple sources +- Task requires fact-checking or verification + +## Workflow + +### 1. Define Scope +Before searching, clarify: +- What specific questions need answering? +- What depth of coverage is needed? +- Are there preferred or authoritative sources? + +### 2. Search Strategy +Execute searches in this order: +1. Broad overview search +2. Specific fact-finding searches +3. Alternative viewpoint searches + +### 3. Source Evaluation +For each source, assess: +- Authority (who wrote it?) +- Currency (when published?) +- Relevance (does it answer the question?) + +### 4. Synthesis +Combine findings into a coherent summary with citations. + +## Examples + +### Example: Technology Research +User: "Research the latest developments in WebAssembly" + +Step 1: Define scope +- Focus on recent developments +- Include browser support status +- Cover WASI developments + +Step 2: Execute searches +- "WebAssembly developments 2026" +- "WASI preview 2 status" +- "WebAssembly component model" + +Step 3: Evaluate sources +- Prefer official documentation +- Cross-reference with reputable tech blogs + +Step 4: Synthesize findings +- Create summary with citations +- Note any conflicting information diff --git a/examples/workflows/harness-pipeline.workflow.yaml b/examples/workflows/harness-pipeline.workflow.yaml new file mode 100644 index 00000000..34346966 --- /dev/null +++ b/examples/workflows/harness-pipeline.workflow.yaml @@ -0,0 +1,73 @@ +# Harness Pipeline Workflow +# Demonstrates: Multi-step workflow using agentic harness agents +# +# This workflow uses agents configured with harness: agentic in their dotprompt. +# Each agent runs through the AgenticExecutor with doom loop detection, +# context compaction, and tool hooks. + +id: harness-pipeline +version: "1.0" +name: "Harness Agent Pipeline" +description: "Demonstrates workflow using agentic harness agents" + +inputSchema: + type: object + properties: + project_name: + type: string + description: "Name of the project to analyze" + analysis_type: + type: string + enum: [quick, thorough] + default: quick + required: [project_name] + +outputSchema: + type: object + properties: + summary: + type: string + files_created: + type: array + items: + type: string + analysis_complete: + type: boolean + +start: setup + +states: + # Step 1: Setup agent creates project structure + - name: setup + type: agent + agent: "harness-test" # Uses harness: agentic in dotprompt + input: + task: | + Create a directory called '{{workflow.input.project_name}}' and inside it: + 1. Create a README.md with project name and description + 2. Create a config.json with version "1.0" + resultPath: "steps.setup" + next: verify + + # Step 2: Verify agent checks the created files + - name: verify + type: agent + agent: "harness-test" # Uses harness: agentic in dotprompt + input: + task: | + List all files in the '{{workflow.input.project_name}}' directory and + read the config.json to confirm version is "1.0". + Report what you found. + resultPath: "steps.verify" + end: true + +# Example trigger: +# curl -X POST http://localhost:8585/api/v1/workflow-runs \ +# -H "Content-Type: application/json" \ +# -d '{ +# "workflow_id": "harness-pipeline", +# "input": { +# "project_name": "my-test-project", +# "analysis_type": "quick" +# } +# }' diff --git a/go.mod b/go.mod index 37d23452..fec9b008 100644 --- a/go.mod +++ b/go.mod @@ -8,6 +8,7 @@ require ( cloud.google.com/go v0.120.0 // indirect cloud.google.com/go/auth v0.16.2 // indirect cloud.google.com/go/compute/metadata v0.7.0 // indirect + connectrpc.com/connect v1.19.1 // indirect github.com/99designs/gqlgen v0.17.75 // indirect github.com/Khan/genqlient v0.8.1 // indirect github.com/KyleBanks/depth v1.2.1 // indirect diff --git a/go.sum b/go.sum index f46b3956..4a0a7fc9 100644 --- a/go.sum +++ b/go.sum @@ -4,6 +4,8 @@ cloud.google.com/go/auth v0.16.2 h1:QvBAGFPLrDeoiNjyfVunhQ10HKNYuOwZ5noee0M5df4= cloud.google.com/go/auth v0.16.2/go.mod h1:sRBas2Y1fB1vZTdurouM0AzuYQBMZinrUYL8EufhtEA= cloud.google.com/go/compute/metadata v0.7.0 h1:PBWF+iiAerVNe8UCHxdOt6eHLVc3ydFeOCw78U8ytSU= cloud.google.com/go/compute/metadata v0.7.0/go.mod h1:j5MvL9PprKL39t166CoB1uVHfQMs4tFQZZcKwksXUjo= +connectrpc.com/connect v1.19.1 h1:R5M57z05+90EfEvCY1b7hBxDVOUl45PrtXtAV2fOC14= +connectrpc.com/connect v1.19.1/go.mod h1:tN20fjdGlewnSFeZxLKb0xwIZ6ozc3OQs2hTXy4du9w= dagger.io/dagger v0.18.16 h1:kgrkJ6tMbqnOie1rYjaquEP5sWqqjAUls0b7oUZnMXE= dagger.io/dagger v0.18.16/go.mod h1:azlZ24m2br95t0jQHUBpL5SiafeqtVDLl1Itlq6GO+4= github.com/99designs/gqlgen v0.17.75 h1:GwHJsptXWLHeY7JO8b7YueUI4w9Pom6wJTICosDtQuI= diff --git a/internal/config/config.go b/internal/config/config.go index b32869ec..40d419c5 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -58,6 +58,8 @@ type Config struct { Secrets SecretsConfig // Faker Templates (for local development) FakerTemplates map[string]FakerTemplate + // Harness Configuration (agentic execution harness) + Harness HarnessConfig // Note: Station now uses official GenKit v1.0.1 plugins (custom plugin code preserved) // Legacy fields (deprecated, use Telemetry struct instead) @@ -74,6 +76,46 @@ type FakerTemplate struct { Model string `yaml:"model"` } +type HarnessConfig struct { + Workspace HarnessWorkspaceConfig `yaml:"workspace"` + Compaction HarnessCompactionConfig `yaml:"compaction"` + Git HarnessGitConfig `yaml:"git"` + NATS HarnessNATSConfig `yaml:"nats"` + Permissions HarnessPermissionsConfig `yaml:"permissions"` +} + +type HarnessWorkspaceConfig struct { + Path string `yaml:"path"` + Mode string `yaml:"mode"` +} + +type HarnessCompactionConfig struct { + Enabled bool `yaml:"enabled"` + Threshold float64 `yaml:"threshold"` + ProtectTokens int `yaml:"protect_tokens"` +} + +type HarnessGitConfig struct { + AutoBranch bool `yaml:"auto_branch"` + BranchPrefix string `yaml:"branch_prefix"` + AutoCommit bool `yaml:"auto_commit"` + RequireApproval bool `yaml:"require_approval"` + WorkflowBranchStrategy string `yaml:"workflow_branch_strategy"` +} + +type HarnessNATSConfig struct { + Enabled bool `yaml:"enabled"` + KVBucket string `yaml:"kv_bucket"` + ObjectBucket string `yaml:"object_bucket"` + MaxFileSize string `yaml:"max_file_size"` + TTL string `yaml:"ttl"` +} + +type HarnessPermissionsConfig struct { + ExternalDirectory string `yaml:"external_directory"` + Bash map[string]string `yaml:"bash"` +} + // CloudShipConfig holds CloudShip Lighthouse integration settings type CloudShipConfig struct { Enabled bool `yaml:"enabled"` // Enable CloudShip integration @@ -633,6 +675,21 @@ func bindEnvVars() { viper.BindEnv("telemetry.sample_rate", "STN_TELEMETRY_SAMPLE_RATE") viper.BindEnv("otel_endpoint", "OTEL_EXPORTER_OTLP_ENDPOINT", "STN_OTEL_ENDPOINT") viper.BindEnv("jaeger_query_url", "STN_JAEGER_QUERY_URL", "JAEGER_QUERY_URL") + + // Harness config + viper.BindEnv("harness.workspace.path", "STN_HARNESS_WORKSPACE_PATH") + viper.BindEnv("harness.workspace.mode", "STN_HARNESS_WORKSPACE_MODE") + viper.BindEnv("harness.compaction.enabled", "STN_HARNESS_COMPACTION_ENABLED") + viper.BindEnv("harness.compaction.threshold", "STN_HARNESS_COMPACTION_THRESHOLD") + viper.BindEnv("harness.compaction.protect_tokens", "STN_HARNESS_COMPACTION_PROTECT_TOKENS") + viper.BindEnv("harness.git.auto_branch", "STN_HARNESS_GIT_AUTO_BRANCH") + viper.BindEnv("harness.git.branch_prefix", "STN_HARNESS_GIT_BRANCH_PREFIX") + viper.BindEnv("harness.git.auto_commit", "STN_HARNESS_GIT_AUTO_COMMIT") + viper.BindEnv("harness.git.require_approval", "STN_HARNESS_GIT_REQUIRE_APPROVAL") + viper.BindEnv("harness.nats.enabled", "STN_HARNESS_NATS_ENABLED") + viper.BindEnv("harness.nats.kv_bucket", "STN_HARNESS_NATS_KV_BUCKET") + viper.BindEnv("harness.nats.object_bucket", "STN_HARNESS_NATS_OBJECT_BUCKET") + viper.BindEnv("harness.permissions.external_directory", "STN_HARNESS_EXTERNAL_DIRECTORY") } func Load() (*Config, error) { @@ -755,6 +812,43 @@ func Load() (*Config, error) { }, }, }, + Harness: HarnessConfig{ + Workspace: HarnessWorkspaceConfig{ + Path: "./workspace", + Mode: "host", + }, + Compaction: HarnessCompactionConfig{ + Enabled: true, + Threshold: 0.85, + ProtectTokens: 40000, + }, + Git: HarnessGitConfig{ + AutoBranch: true, + BranchPrefix: "agent/", + AutoCommit: false, + RequireApproval: true, + WorkflowBranchStrategy: "shared", + }, + NATS: HarnessNATSConfig{ + Enabled: true, + KVBucket: "harness-state", + ObjectBucket: "harness-files", + MaxFileSize: "100MB", + TTL: "24h", + }, + Permissions: HarnessPermissionsConfig{ + ExternalDirectory: "deny", + Bash: map[string]string{ + "*": "allow", + "rm -rf *": "deny", + "rm -r *": "deny", + "git push --force*": "deny", + "git reset --hard*": "deny", + "git push *": "ask", + "git commit *": "ask", + }, + }, + }, // Legacy fields for backward compatibility TelemetryEnabled: true, OTELEndpoint: LocalJaegerEndpoint, @@ -1122,6 +1216,58 @@ func Load() (*Config, error) { cfg.Coding.NATS.ObjectStore = viper.GetString("coding.nats.object_store") } + if viper.IsSet("harness.workspace.path") { + cfg.Harness.Workspace.Path = viper.GetString("harness.workspace.path") + } + if viper.IsSet("harness.workspace.mode") { + cfg.Harness.Workspace.Mode = viper.GetString("harness.workspace.mode") + } + if viper.IsSet("harness.compaction.enabled") { + cfg.Harness.Compaction.Enabled = viper.GetBool("harness.compaction.enabled") + } + if viper.IsSet("harness.compaction.threshold") { + cfg.Harness.Compaction.Threshold = viper.GetFloat64("harness.compaction.threshold") + } + if viper.IsSet("harness.compaction.protect_tokens") { + cfg.Harness.Compaction.ProtectTokens = viper.GetInt("harness.compaction.protect_tokens") + } + if viper.IsSet("harness.git.auto_branch") { + cfg.Harness.Git.AutoBranch = viper.GetBool("harness.git.auto_branch") + } + if viper.IsSet("harness.git.branch_prefix") { + cfg.Harness.Git.BranchPrefix = viper.GetString("harness.git.branch_prefix") + } + if viper.IsSet("harness.git.auto_commit") { + cfg.Harness.Git.AutoCommit = viper.GetBool("harness.git.auto_commit") + } + if viper.IsSet("harness.git.require_approval") { + cfg.Harness.Git.RequireApproval = viper.GetBool("harness.git.require_approval") + } + if viper.IsSet("harness.git.workflow_branch_strategy") { + cfg.Harness.Git.WorkflowBranchStrategy = viper.GetString("harness.git.workflow_branch_strategy") + } + if viper.IsSet("harness.nats.enabled") { + cfg.Harness.NATS.Enabled = viper.GetBool("harness.nats.enabled") + } + if viper.IsSet("harness.nats.kv_bucket") { + cfg.Harness.NATS.KVBucket = viper.GetString("harness.nats.kv_bucket") + } + if viper.IsSet("harness.nats.object_bucket") { + cfg.Harness.NATS.ObjectBucket = viper.GetString("harness.nats.object_bucket") + } + if viper.IsSet("harness.nats.max_file_size") { + cfg.Harness.NATS.MaxFileSize = viper.GetString("harness.nats.max_file_size") + } + if viper.IsSet("harness.nats.ttl") { + cfg.Harness.NATS.TTL = viper.GetString("harness.nats.ttl") + } + if viper.IsSet("harness.permissions.external_directory") { + cfg.Harness.Permissions.ExternalDirectory = viper.GetString("harness.permissions.external_directory") + } + if viper.IsSet("harness.permissions.bash") { + cfg.Harness.Permissions.Bash = viper.GetStringMapString("harness.permissions.bash") + } + // Load faker templates from config file cfg.FakerTemplates = loadFakerTemplates() diff --git a/internal/config/schema.go b/internal/config/schema.go index 9a82df35..65f42ac2 100644 --- a/internal/config/schema.go +++ b/internal/config/schema.go @@ -36,13 +36,14 @@ type ConfigSection struct { var ConfigSections = []ConfigSection{ {Name: "ai", Description: "AI Provider Settings", Order: 1}, {Name: "coding", Description: "Coding Backend Settings", Order: 2}, - {Name: "cloudship", Description: "CloudShip Integration", Order: 3}, - {Name: "lattice", Description: "Lattice Mesh Network", Order: 4}, - {Name: "telemetry", Description: "Telemetry & Observability", Order: 5}, - {Name: "sandbox", Description: "Sandbox Execution", Order: 6}, - {Name: "webhook", Description: "Webhook Settings", Order: 7}, - {Name: "notifications", Description: "Notification Settings", Order: 8}, - {Name: "server", Description: "Server & Port Settings", Order: 9}, + {Name: "harness", Description: "Agentic Harness Settings", Order: 3}, + {Name: "cloudship", Description: "CloudShip Integration", Order: 4}, + {Name: "lattice", Description: "Lattice Mesh Network", Order: 5}, + {Name: "telemetry", Description: "Telemetry & Observability", Order: 6}, + {Name: "sandbox", Description: "Sandbox Execution", Order: 7}, + {Name: "webhook", Description: "Webhook Settings", Order: 8}, + {Name: "notifications", Description: "Notification Settings", Order: 9}, + {Name: "server", Description: "Server & Port Settings", Order: 10}, } var ConfigSchema = []ConfigField{ @@ -84,6 +85,32 @@ var ConfigSchema = []ConfigField{ {Key: "coding.git.user_name", Type: FieldTypeString, Description: "Git commit author name", Default: "Station Bot", Section: "coding"}, {Key: "coding.git.user_email", Type: FieldTypeString, Description: "Git commit author email", Default: "station@cloudship.ai", Section: "coding"}, + // Harness Workspace Settings + {Key: "harness.workspace.path", Type: FieldTypeString, Description: "Workspace directory path (relative to config)", Default: "./workspace", Section: "harness"}, + {Key: "harness.workspace.mode", Type: FieldTypeString, Description: "Workspace mode (host, sandbox, agentfs)", Default: "host", Section: "harness", Options: []string{"host", "sandbox", "agentfs"}}, + + // Harness Compaction Settings + {Key: "harness.compaction.enabled", Type: FieldTypeBool, Description: "Enable automatic context compaction", Default: true, Section: "harness"}, + {Key: "harness.compaction.threshold", Type: FieldTypeString, Description: "Context window threshold for compaction (0.0-1.0)", Default: "0.85", Section: "harness"}, + {Key: "harness.compaction.protect_tokens", Type: FieldTypeInt, Description: "Recent tokens to protect from pruning", Default: 40000, Section: "harness"}, + + // Harness Git Settings + {Key: "harness.git.auto_branch", Type: FieldTypeBool, Description: "Auto-create branch for each task", Default: true, Section: "harness"}, + {Key: "harness.git.branch_prefix", Type: FieldTypeString, Description: "Prefix for agent-created branches", Default: "agent/", Section: "harness"}, + {Key: "harness.git.auto_commit", Type: FieldTypeBool, Description: "Auto-commit on task completion", Default: false, Section: "harness"}, + {Key: "harness.git.require_approval", Type: FieldTypeBool, Description: "Require approval before push/PR", Default: true, Section: "harness"}, + {Key: "harness.git.workflow_branch_strategy", Type: FieldTypeString, Description: "Branch strategy in workflows (shared, per_agent)", Default: "shared", Section: "harness", Options: []string{"shared", "per_agent"}}, + + // Harness NATS Settings + {Key: "harness.nats.enabled", Type: FieldTypeBool, Description: "Enable NATS for state sharing", Default: true, Section: "harness"}, + {Key: "harness.nats.kv_bucket", Type: FieldTypeString, Description: "NATS KV bucket for state", Default: "harness-state", Section: "harness"}, + {Key: "harness.nats.object_bucket", Type: FieldTypeString, Description: "NATS Object Store bucket", Default: "harness-files", Section: "harness"}, + {Key: "harness.nats.max_file_size", Type: FieldTypeString, Description: "Max file size for object store", Default: "100MB", Section: "harness"}, + {Key: "harness.nats.ttl", Type: FieldTypeString, Description: "TTL for stored artifacts", Default: "24h", Section: "harness"}, + + // Harness Permission Settings + {Key: "harness.permissions.external_directory", Type: FieldTypeString, Description: "Access to paths outside workspace (deny, ask, allow)", Default: "deny", Section: "harness", Options: []string{"deny", "ask", "allow"}}, + // CloudShip Integration {Key: "cloudship.enabled", Type: FieldTypeBool, Description: "Enable CloudShip integration", Default: false, Section: "cloudship"}, {Key: "cloudship.api_key", Type: FieldTypeString, Description: "CloudShip personal API key (cst_...)", Secret: true, Section: "cloudship"}, diff --git a/internal/mcp/agent_handlers.go b/internal/mcp/agent_handlers.go index 4bccb680..f64c352a 100644 --- a/internal/mcp/agent_handlers.go +++ b/internal/mcp/agent_handlers.go @@ -122,6 +122,15 @@ func (s *Server) handleCreateAgent(ctx context.Context, request mcp.CallToolRequ } notifyEnabled := request.GetBool("notify", false) + + harnessConfig := request.GetString("harness_config", "") + if harnessConfig != "" { + var harnessMap map[string]interface{} + if err := json.Unmarshal([]byte(harnessConfig), &harnessMap); err != nil { + return mcp.NewToolResultError(fmt.Sprintf("Invalid harness_config JSON: %v. Example: {\"max_steps\": 50, \"sandbox\": {\"mode\": \"docker\"}}", err)), nil + } + } + var toolNames []string if request.Params.Arguments != nil { if argsMap, ok := request.Params.Arguments.(map[string]interface{}); ok { @@ -173,7 +182,7 @@ func (s *Server) handleCreateAgent(ctx context.Context, request mcp.CallToolRequ } if s.agentExportService != nil { - if err := s.agentExportService.ExportAgentWithConfigs(createdAgent.ID, app, appType, sandboxConfig, codingConfig, notifyEnabled); err != nil { + if err := s.agentExportService.ExportAgentWithAllConfigs(createdAgent.ID, app, appType, sandboxConfig, codingConfig, harnessConfig, notifyEnabled); err != nil { response["export_warning"] = fmt.Sprintf("Agent created but export failed: %v. Use 'stn agent export %s' to export manually.", err, name) } } @@ -351,6 +360,14 @@ func (s *Server) handleUpdateAgent(ctx context.Context, request mcp.CallToolRequ } } + harnessConfig := request.GetString("harness_config", "") + if harnessConfig != "" && harnessConfig != "{}" { + var harnessMap map[string]interface{} + if err := json.Unmarshal([]byte(harnessConfig), &harnessMap); err != nil { + return mcp.NewToolResultError(fmt.Sprintf("Invalid harness_config JSON: %v. Example: {\"max_steps\": 50, \"sandbox\": {\"mode\": \"docker\"}}", err)), nil + } + } + notifyEnabled := request.GetBool("notify", false) err = s.repos.Agents.UpdateWithMemory( @@ -375,7 +392,7 @@ func (s *Server) handleUpdateAgent(ctx context.Context, request mcp.CallToolRequ } if s.agentExportService != nil { - if err := s.agentExportService.ExportAgentWithConfigs(agentID, app, appType, sandboxConfig, codingConfig, notifyEnabled); err != nil { + if err := s.agentExportService.ExportAgentWithAllConfigs(agentID, app, appType, sandboxConfig, codingConfig, harnessConfig, notifyEnabled); err != nil { response := map[string]interface{}{ "success": true, "agent": map[string]interface{}{ diff --git a/internal/mcp/tools_setup.go b/internal/mcp/tools_setup.go index c3b0e200..03ddf7cb 100644 --- a/internal/mcp/tools_setup.go +++ b/internal/mcp/tools_setup.go @@ -10,7 +10,7 @@ import ( func (s *Server) setupTools() { // Agent management tools (CRUD operations) createAgentTool := mcp.NewTool("create_agent", - mcp.WithDescription("Create a new AI agent with specified configuration. For sandbox options, read 'station://docs/sandbox' resource."), + mcp.WithDescription("Create a new AI agent with specified configuration. For sandbox options, read 'station://docs/sandbox' resource. For agentic harness (Claude Agent SDK-like execution), use harness_config."), mcp.WithString("name", mcp.Required(), mcp.Description("Name of the agent")), mcp.WithString("description", mcp.Required(), mcp.Description("Description of what the agent does")), mcp.WithString("prompt", mcp.Required(), mcp.Description("System prompt for the agent")), @@ -28,12 +28,13 @@ func (s *Server) setupTools() { mcp.WithString("sandbox", mcp.Description("Sandbox configuration as JSON. Read 'station://docs/sandbox' for all options. Modes: 'compute' (single execution, default) or 'code' (persistent session). For code mode, use 'session': 'workflow' to share sandbox across workflow steps, or 'agent' (default) for per-agent sessions. Example: {\"mode\": \"code\", \"session\": \"workflow\", \"runtime\": \"python\", \"pip_packages\": [\"pandas\"]}. Full options: mode, session, runtime, image, timeout_seconds, allow_network, pip_packages, npm_packages, limits.")), mcp.WithString("coding", mcp.Description("AI coding backend configuration as JSON. Enables coding_open, code, coding_close, coding_commit, coding_push tools. Backends: 'claudecode' (Claude Code CLI), 'opencode-cli' (OpenCode CLI), 'opencode' (OpenCode HTTP), 'opencode-nats' (NATS). Example: {\"enabled\": true, \"backend\": \"claudecode\"}. Options: enabled (bool), backend (string), workspace_path (optional).")), mcp.WithBoolean("notify", mcp.Description("Enable the notify tool for this agent to send webhook notifications (ntfy, Slack, Discord, etc.). Requires notify webhook to be configured in Station config. Default: false.")), + mcp.WithString("harness_config", mcp.Description("Agentic harness configuration as JSON. Enables Claude Agent SDK-like multi-turn execution with sandbox isolation. Example: {\"max_steps\": 50, \"doom_loop_threshold\": 5, \"timeout\": \"30m\", \"sandbox\": {\"mode\": \"docker\", \"image\": \"python:3.11\"}}. Options: max_steps (int, default 25), doom_loop_threshold (int, default 3), timeout (duration string), sandbox (object with mode: host|docker|e2b). When set, adds 'harness: agentic' to frontmatter. Tools should include: read, write, bash, glob, grep, edit.")), ) s.mcpServer.AddTool(createAgentTool, s.handleCreateAgent) // Update agent tool updateAgentTool := mcp.NewTool("update_agent", - mcp.WithDescription("Update an existing agent's configuration. Note: To manage tools, use add_tool/remove_tool instead. For sandbox options, read 'station://docs/sandbox' resource."), + mcp.WithDescription("Update an existing agent's configuration. Note: To manage tools, use add_tool/remove_tool instead. For sandbox options, read 'station://docs/sandbox' resource. For agentic harness, use harness_config."), mcp.WithString("agent_id", mcp.Required(), mcp.Description("ID of the agent to update")), mcp.WithString("name", mcp.Description("New name for the agent")), mcp.WithString("description", mcp.Description("New description for the agent")), @@ -47,6 +48,7 @@ func (s *Server) setupTools() { mcp.WithString("sandbox", mcp.Description("Sandbox configuration as JSON. Read 'station://docs/sandbox' for all options. For code mode with workflow scoping: {\"mode\": \"code\", \"session\": \"workflow\", \"runtime\": \"python\"}. Set to \"{}\" to remove sandbox. Options: mode (compute/code), session (workflow/agent), runtime, pip_packages, npm_packages, timeout_seconds, allow_network, limits.")), mcp.WithString("coding", mcp.Description("AI coding backend configuration as JSON. Set to \"{}\" to remove. Backends: 'claudecode' (Claude Code CLI), 'opencode-cli', 'opencode', 'opencode-nats'. Example: {\"enabled\": true, \"backend\": \"claudecode\"}.")), mcp.WithBoolean("notify", mcp.Description("Enable/disable the notify tool for this agent. Set to true to enable webhook notifications, false to disable.")), + mcp.WithString("harness_config", mcp.Description("Agentic harness configuration as JSON. Set to \"{}\" to remove. Example: {\"max_steps\": 50, \"sandbox\": {\"mode\": \"docker\"}}. Options: max_steps, doom_loop_threshold, timeout, sandbox (mode: host|docker|e2b).")), ) s.mcpServer.AddTool(updateAgentTool, s.handleUpdateAgent) diff --git a/internal/services/agent_execution_engine.go b/internal/services/agent_execution_engine.go index a87a7ba5..0e62b0be 100644 --- a/internal/services/agent_execution_engine.go +++ b/internal/services/agent_execution_engine.go @@ -14,11 +14,18 @@ import ( "station/internal/logging" "station/internal/storage" dotprompt "station/pkg/dotprompt" + "station/pkg/harness" + "station/pkg/harness/git" + "station/pkg/harness/sandbox" + "station/pkg/harness/stream" + "station/pkg/harness/tools" + "station/pkg/harness/workspace" "station/pkg/models" "station/pkg/schema" "station/pkg/types" "github.com/firebase/genkit/go/ai" + "github.com/firebase/genkit/go/genkit" "github.com/firebase/genkit/go/plugins/mcp" "github.com/google/uuid" "go.opentelemetry.io/otel/attribute" @@ -520,6 +527,11 @@ func (aee *AgentExecutionEngine) ExecuteWithOptions(ctx context.Context, agent * } } + harnessMode := aee.parseHarnessFromAgent(agent, environment.Name) + if harnessMode == "agentic" { + logging.Info("Agent %s configured with agentic harness mode", agent.Name) + } + sandboxConfig := aee.parseSandboxConfigFromAgent(agent, environment.Name) toolExecCtx := ExecutionContext{ WorkflowRunID: workflowRunID, @@ -618,14 +630,19 @@ func (aee *AgentExecutionEngine) ExecuteWithOptions(ctx context.Context, agent * } } - // Use clean, unified dotprompt.Execute() execution path with context for timeout protection - logging.Info("[TOOLS DEBUG] About to call ExecuteAgent with %d mcpTools for agent %s", len(mcpTools), agent.Name) - for i, t := range mcpTools { - if tool, ok := t.(interface{ Name() string }); ok { - logging.Info("[TOOLS DEBUG] Tool[%d]: %s", i, tool.Name()) + var response *dotprompt.ExecutionResponse + + if harnessMode == "agentic" { + response, err = aee.executeWithAgenticHarness(ctx, agent, genkitApp, mcpTools, task, environment.Name, runUUID, workflowRunID) + } else { + logging.Info("[TOOLS DEBUG] About to call ExecuteAgent with %d mcpTools for agent %s", len(mcpTools), agent.Name) + for i, t := range mcpTools { + if tool, ok := t.(interface{ Name() string }); ok { + logging.Info("[TOOLS DEBUG] Tool[%d]: %s", i, tool.Name()) + } } + response, err = executor.ExecuteAgent(ctx, *agent, agentTools, genkitApp, mcpTools, task, logCallback, environment.Name, userVariables) } - response, err := executor.ExecuteAgent(ctx, *agent, agentTools, genkitApp, mcpTools, task, logCallback, environment.Name, userVariables) // Clean up MCP connections after execution is complete aee.mcpConnManager.CleanupConnections(aee.activeMCPClients) @@ -1151,6 +1168,26 @@ type latticeFrontmatter struct { Lattice bool `yaml:"lattice"` } +type harnessFrontmatter struct { + Harness string `yaml:"harness"` + HarnessConfig *agentHarnessFrontmatter `yaml:"harness_config"` +} + +type agentHarnessFrontmatter struct { + MaxSteps int `yaml:"max_steps"` + DoomLoopThreshold int `yaml:"doom_loop_threshold"` + Timeout string `yaml:"timeout"` + Sandbox *harnessSandboxFrontmatter `yaml:"sandbox"` +} + +type harnessSandboxFrontmatter struct { + Mode string `yaml:"mode"` + Image string `yaml:"image"` + Timeout string `yaml:"timeout"` + Network bool `yaml:"network"` + Environment map[string]string `yaml:"environment"` +} + func (aee *AgentExecutionEngine) parseLatticeEnabledFromAgent(agent *models.Agent) bool { if agent == nil { return false @@ -1279,3 +1316,299 @@ func (aee *AgentExecutionEngine) parseCodingConfigFromAgent(agent *models.Agent, return fm.Coding } + +func (aee *AgentExecutionEngine) parseHarnessFromAgent(agent *models.Agent, environmentName string) string { + fm := aee.parseFullHarnessFrontmatter(agent, environmentName) + if fm == nil { + return "" + } + return fm.Harness +} + +func (aee *AgentExecutionEngine) parseFullHarnessFrontmatter(agent *models.Agent, environmentName string) *harnessFrontmatter { + promptPath := config.GetAgentPromptPath(environmentName, agent.Name) + + content, err := os.ReadFile(promptPath) + if err != nil { + logging.Debug("Failed to read dotprompt file for harness config: %v", err) + return nil + } + + parts := strings.Split(string(content), "---") + if len(parts) < 3 { + return nil + } + + yamlContent := strings.TrimSpace(parts[1]) + if yamlContent == "" { + return nil + } + + var fm harnessFrontmatter + if err := yaml.Unmarshal([]byte(yamlContent), &fm); err != nil { + logging.Debug("Failed to parse harness config from dotprompt: %v", err) + return nil + } + + return &fm +} + +func (aee *AgentExecutionEngine) buildAgentHarnessConfig(fm *harnessFrontmatter) *harness.AgentHarnessConfig { + cfg := harness.DefaultAgentHarnessConfig() + + if fm == nil || fm.HarnessConfig == nil { + return cfg + } + + if fm.HarnessConfig.MaxSteps > 0 { + cfg.MaxSteps = fm.HarnessConfig.MaxSteps + } + if fm.HarnessConfig.DoomLoopThreshold > 0 { + cfg.DoomLoopThreshold = fm.HarnessConfig.DoomLoopThreshold + } + if fm.HarnessConfig.Timeout != "" { + if d, err := time.ParseDuration(fm.HarnessConfig.Timeout); err == nil { + cfg.Timeout = d + } + } + + return cfg +} + +// convertConfigToHarnessConfig converts internal config.HarnessConfig to harness.HarnessConfig +func convertConfigToHarnessConfig(cfg config.HarnessConfig) *harness.HarnessConfig { + return &harness.HarnessConfig{ + Workspace: harness.WorkspaceConfig{ + Path: cfg.Workspace.Path, + Mode: cfg.Workspace.Mode, + }, + Compaction: harness.CompactionConfig{ + Enabled: cfg.Compaction.Enabled, + Threshold: cfg.Compaction.Threshold, + ProtectTokens: cfg.Compaction.ProtectTokens, + }, + Git: harness.GitConfig{ + AutoBranch: cfg.Git.AutoBranch, + BranchPrefix: cfg.Git.BranchPrefix, + AutoCommit: cfg.Git.AutoCommit, + RequireApproval: cfg.Git.RequireApproval, + WorkflowBranchStrategy: cfg.Git.WorkflowBranchStrategy, + }, + NATS: harness.NATSConfig{ + Enabled: cfg.NATS.Enabled, + KVBucket: cfg.NATS.KVBucket, + ObjectBucket: cfg.NATS.ObjectBucket, + MaxFileSize: cfg.NATS.MaxFileSize, + TTL: cfg.NATS.TTL, + }, + Permissions: harness.PermissionsConfig{ + ExternalDirectory: cfg.Permissions.ExternalDirectory, + Bash: cfg.Permissions.Bash, + }, + } +} + +func convertHarnessResultToExecutionResponse(result *harness.ExecutionResult, modelName string) *dotprompt.ExecutionResponse { + return &dotprompt.ExecutionResponse{ + Success: result.Success, + Response: result.Response, + Duration: result.Duration, + ModelName: modelName, + StepsUsed: result.TotalSteps, + TokenUsage: map[string]interface{}{ + "total_tokens": result.TotalTokens, + }, + Error: result.Error, + } +} + +func (aee *AgentExecutionEngine) createHarnessSandbox(ctx context.Context, fm *harnessFrontmatter, workspacePath string) (sandbox.Sandbox, error) { + if fm == nil || fm.HarnessConfig == nil || fm.HarnessConfig.Sandbox == nil { + return nil, nil + } + + sbFM := fm.HarnessConfig.Sandbox + mode := sandbox.Mode(sbFM.Mode) + if mode == "" { + mode = sandbox.ModeHost + } + + cfg := sandbox.Config{ + Mode: mode, + Image: sbFM.Image, + WorkspacePath: workspacePath, + Environment: sbFM.Environment, + Network: sandbox.NetworkConfig{ + Enabled: sbFM.Network, + }, + } + + if sbFM.Timeout != "" { + if d, err := time.ParseDuration(sbFM.Timeout); err == nil { + cfg.Timeout = d + } + } + + factory := sandbox.NewFactory(sandbox.DefaultConfig()) + sb, err := factory.Create(cfg) + if err != nil { + return nil, fmt.Errorf("failed to create sandbox: %w", err) + } + + if err := sb.Create(ctx); err != nil { + return nil, fmt.Errorf("failed to initialize sandbox: %w", err) + } + + logging.Info("Sandbox created: mode=%s, image=%s, workspace=%s", mode, sbFM.Image, workspacePath) + return sb, nil +} + +func (aee *AgentExecutionEngine) executeWithAgenticHarness( + ctx context.Context, + agent *models.Agent, + genkitApp *genkit.Genkit, + mcpTools []ai.ToolRef, + task string, + environmentName string, + runUUID string, + workflowRunID string, +) (*dotprompt.ExecutionResponse, error) { + startTime := time.Now() + + loadedConfig := config.GetLoadedConfig() + harnessConfig := convertConfigToHarnessConfig(loadedConfig.Harness) + + harnessFM := aee.parseFullHarnessFrontmatter(agent, environmentName) + agentHarnessConfig := aee.buildAgentHarnessConfig(harnessFM) + + workspacePath := harnessConfig.Workspace.Path + if workspacePath == "" { + workspacePath = "./workspace" + } + + ws := workspace.NewHostWorkspace(workspacePath) + if err := ws.Initialize(ctx); err != nil { + return &dotprompt.ExecutionResponse{ + Success: false, + Error: fmt.Sprintf("workspace initialization failed: %v", err), + Duration: time.Since(startTime), + }, err + } + + sb, err := aee.createHarnessSandbox(ctx, harnessFM, workspacePath) + if err != nil { + return &dotprompt.ExecutionResponse{ + Success: false, + Error: fmt.Sprintf("sandbox creation failed: %v", err), + Duration: time.Since(startTime), + }, err + } + if sb != nil { + defer func() { + if destroyErr := sb.Destroy(ctx); destroyErr != nil { + logging.Error("Failed to destroy sandbox: %v", destroyErr) + } + }() + } + + toolRegistry := tools.NewToolRegistryWithSandbox(genkitApp, workspacePath, sb) + if err := toolRegistry.RegisterBuiltinTools(); err != nil { + return &dotprompt.ExecutionResponse{ + Success: false, + Error: fmt.Sprintf("tool registration failed: %v", err), + Duration: time.Since(startTime), + }, err + } + + allTools := toolRegistry.All() + + for _, toolRef := range mcpTools { + if tool, ok := toolRef.(ai.Tool); ok { + allTools = append(allTools, tool) + } + } + + modelName := loadedConfig.AIModel + if modelName == "" { + modelName = "openai/gpt-4o-mini" + } + + executorOpts := []harness.ExecutorOption{ + harness.WithWorkspace(ws), + harness.WithModelName(modelName), + } + + if sb != nil { + executorOpts = append(executorOpts, harness.WithSandbox(sb)) + } + + if harnessConfig.Git.AutoBranch { + gitMgr := git.NewManager(workspacePath, git.WithBranchPrefix(harnessConfig.Git.BranchPrefix)) + if gitMgr.IsGitRepo(ctx) { + executorOpts = append(executorOpts, harness.WithGitManager(gitMgr)) + logging.Info("Git integration enabled for agent %s (auto_branch=%v, auto_commit=%v)", + agent.Name, harnessConfig.Git.AutoBranch, harnessConfig.Git.AutoCommit) + } + } + + if harnessConfig.Compaction.Enabled { + contextWindow := 128000 + compactor := harness.NewCompactor(genkitApp, modelName, harnessConfig.Compaction, contextWindow) + executorOpts = append(executorOpts, harness.WithCompactor(compactor)) + logging.Info("Context compaction enabled for agent %s (threshold=%.0f%%, protect=%d tokens)", + agent.Name, harnessConfig.Compaction.Threshold*100, harnessConfig.Compaction.ProtectTokens) + } + + streamPublisher := aee.getStreamPublisher(harnessConfig) + if streamPublisher != nil { + executorOpts = append(executorOpts, harness.WithStreamPublisher(streamPublisher)) + logging.Info("Stream publisher enabled for agent %s", agent.Name) + } + + agenticExecutor := harness.NewAgenticExecutor( + genkitApp, + harnessConfig, + agentHarnessConfig, + executorOpts..., + ) + + logging.Info("Executing agent %s with agentic harness (max_steps=%d, doom_threshold=%d, timeout=%v, tools=%d)", + agent.Name, agentHarnessConfig.MaxSteps, agentHarnessConfig.DoomLoopThreshold, agentHarnessConfig.Timeout, len(allTools)) + + stationID := "" + if loadedConfig.CloudShip.StationID != "" { + stationID = loadedConfig.CloudShip.StationID + } else if loadedConfig.Lattice.StationID != "" { + stationID = loadedConfig.Lattice.StationID + } + + execOpts := harness.ExecuteOptions{ + StationRunID: fmt.Sprintf("%d", agent.ID), + RunUUID: runUUID, + WorkflowRunID: workflowRunID, + SessionID: workflowRunID, + AgentName: agent.Name, + StationID: stationID, + } + result, err := agenticExecutor.ExecuteWithOptions(ctx, fmt.Sprintf("%d", agent.ID), task, allTools, execOpts) + if err != nil { + logging.Error("Agentic harness execution failed: %v", err) + return &dotprompt.ExecutionResponse{ + Success: false, + Error: err.Error(), + Duration: time.Since(startTime), + }, err + } + + logging.Info("Agentic harness execution completed: success=%v, steps=%d, tokens=%d, finish=%s", + result.Success, result.TotalSteps, result.TotalTokens, result.FinishReason) + + return convertHarnessResultToExecutionResponse(result, modelName), nil +} + +func (aee *AgentExecutionEngine) getStreamPublisher(cfg *harness.HarnessConfig) stream.Publisher { + if cfg == nil || !cfg.Streaming.Enabled { + return nil + } + return stream.NewChannelPublisher(100) +} diff --git a/internal/services/agent_export_service.go b/internal/services/agent_export_service.go index 66a3c9c1..e5752501 100644 --- a/internal/services/agent_export_service.go +++ b/internal/services/agent_export_service.go @@ -44,6 +44,10 @@ func (s *AgentExportService) ExportAgentWithSandbox(agentID int64, app, appType, } func (s *AgentExportService) ExportAgentWithConfigs(agentID int64, app, appType, sandboxConfig, codingConfig string, notifyEnabled bool) error { + return s.ExportAgentWithAllConfigs(agentID, app, appType, sandboxConfig, codingConfig, "", notifyEnabled) +} + +func (s *AgentExportService) ExportAgentWithAllConfigs(agentID int64, app, appType, sandboxConfig, codingConfig, harnessConfig string, notifyEnabled bool) error { agent, err := s.repos.Agents.GetByID(agentID) if err != nil { return fmt.Errorf("failed to get agent: %v", err) @@ -78,7 +82,7 @@ func (s *AgentExportService) ExportAgentWithConfigs(agentID int64, app, appType, return fmt.Errorf("failed to get child agents: %v", err) } - dotpromptContent := s.generateDotpromptContentWithConfigs(agent, toolsWithDetails, childAgents, environment.Name, app, appType, sandboxConfig, codingConfig, notifyEnabled) + dotpromptContent := s.generateDotpromptContentWithAllConfigs(agent, toolsWithDetails, childAgents, environment.Name, app, appType, sandboxConfig, codingConfig, harnessConfig, notifyEnabled) outputPath := config.GetAgentPromptPath(environment.Name, agent.Name) @@ -104,6 +108,10 @@ func (s *AgentExportService) generateDotpromptContentWithSandbox(agent *models.A } func (s *AgentExportService) generateDotpromptContentWithConfigs(agent *models.Agent, tools []*models.AgentToolWithDetails, childAgents []*models.ChildAgent, environmentName, app, appType, sandboxConfig, codingConfig string, notifyEnabled bool) string { + return s.generateDotpromptContentWithAllConfigs(agent, tools, childAgents, environmentName, app, appType, sandboxConfig, codingConfig, "", notifyEnabled) +} + +func (s *AgentExportService) generateDotpromptContentWithAllConfigs(agent *models.Agent, tools []*models.AgentToolWithDetails, childAgents []*models.ChildAgent, environmentName, app, appType, sandboxConfig, codingConfig, harnessConfig string, notifyEnabled bool) string { var toolNames []string for _, tool := range tools { toolNames = append(toolNames, tool.ToolName) @@ -152,6 +160,14 @@ metadata: content += "\nnotify: true" } + if harnessConfig != "" { + harnessYAML := s.convertHarnessConfigJSONToYAML(harnessConfig) + if harnessYAML != "" { + content += "\nharness: agentic" + content += "\nharness_config:\n" + s.indentLines(harnessYAML, " ") + } + } + scheduleSection := s.generateScheduleSection(agent) if scheduleSection != "" { content += scheduleSection @@ -252,6 +268,30 @@ func (s *AgentExportService) convertCodingJSONToYAML(codingJSON string) string { return strings.TrimSpace(string(yamlBytes)) } +func (s *AgentExportService) convertHarnessConfigJSONToYAML(harnessJSON string) string { + if harnessJSON == "" || harnessJSON == "{}" { + return "" + } + + var harnessMap map[string]interface{} + if err := json.Unmarshal([]byte(harnessJSON), &harnessMap); err != nil { + log.Printf("Warning: Failed to parse harness_config JSON: %v", err) + return "" + } + + if len(harnessMap) == 0 { + return "" + } + + yamlBytes, err := yaml.Marshal(harnessMap) + if err != nil { + log.Printf("Warning: Failed to convert harness_config to YAML: %v", err) + return "" + } + + return strings.TrimSpace(string(yamlBytes)) +} + // generateInputSchemaSection generates the input schema for dotprompt func (s *AgentExportService) generateInputSchemaSection(agent *models.Agent) (string, error) { // Use the existing ExportHelper for proper schema generation diff --git a/pkg/harness/compaction.go b/pkg/harness/compaction.go new file mode 100644 index 00000000..a093038b --- /dev/null +++ b/pkg/harness/compaction.go @@ -0,0 +1,428 @@ +package harness + +import ( + "context" + "encoding/json" + "fmt" + "strings" + "time" + + "github.com/firebase/genkit/go/ai" + "github.com/firebase/genkit/go/genkit" +) + +// HistoryBackend interface for persisting conversation history +type HistoryBackend interface { + WriteFile(path string, content []byte) error + ReadFile(path string) ([]byte, error) +} + +type Compactor struct { + genkitApp *genkit.Genkit + modelName string + config CompactionConfig + tokenCounter TokenCounter + contextWindow int + historyBackend HistoryBackend // For history offload + historyPrefix string // e.g., "/conversation_history" + sessionID string // For unique history file names +} + +type TokenCounter interface { + CountTokens(messages []*ai.Message) (int, error) +} + +type SimpleTokenCounter struct{} + +func (s *SimpleTokenCounter) CountTokens(messages []*ai.Message) (int, error) { + total := 0 + for _, msg := range messages { + for _, part := range msg.Content { + if part.IsText() { + total += len(strings.Fields(part.Text)) * 4 / 3 + } + } + } + return total, nil +} + +func NewCompactor(genkitApp *genkit.Genkit, modelName string, config CompactionConfig, contextWindow int) *Compactor { + return &Compactor{ + genkitApp: genkitApp, + modelName: modelName, + config: config, + tokenCounter: &SimpleTokenCounter{}, + contextWindow: contextWindow, + } +} + +func (c *Compactor) WithTokenCounter(counter TokenCounter) *Compactor { + c.tokenCounter = counter + return c +} + +func (c *Compactor) WithHistoryBackend(backend HistoryBackend, prefix, sessionID string) *Compactor { + c.historyBackend = backend + c.historyPrefix = prefix + c.sessionID = sessionID + return c +} + +func (c *Compactor) CompactIfNeeded(ctx context.Context, history []*ai.Message) ([]*ai.Message, bool, error) { + if !c.config.Enabled { + return history, false, nil + } + + currentTokens, err := c.tokenCounter.CountTokens(history) + if err != nil { + return history, false, fmt.Errorf("failed to count tokens: %w", err) + } + + threshold := int(float64(c.contextWindow) * c.config.Threshold) + if currentTokens < threshold { + return history, false, nil + } + + compacted, err := c.compact(ctx, history) + if err != nil { + return history, false, fmt.Errorf("compaction failed: %w", err) + } + + return compacted, true, nil +} + +func (c *Compactor) compact(ctx context.Context, history []*ai.Message) ([]*ai.Message, error) { + if len(history) <= 2 { + return history, nil + } + + protectedTokens := c.config.ProtectTokens + if protectedTokens <= 0 { + protectedTokens = 40000 + } + + currentTokens, _ := c.tokenCounter.CountTokens(history) + tokensToPreserve := 0 + preserveFromIdx := len(history) + + for i := len(history) - 1; i >= 0; i-- { + msgTokens, _ := c.tokenCounter.CountTokens([]*ai.Message{history[i]}) + if tokensToPreserve+msgTokens > protectedTokens { + preserveFromIdx = i + 1 + break + } + tokensToPreserve += msgTokens + } + + if preserveFromIdx <= 1 { + return history, nil + } + + toSummarize := history[1:preserveFromIdx] + if len(toSummarize) == 0 { + return history, nil + } + + // Offload history before summarization (preserves full data) + var historyFile string + if c.config.HistoryOffload { + var err error + historyFile, err = c.offloadHistory(toSummarize) + if err != nil { + // Log but don't fail - offload is optional + fmt.Printf("Warning: failed to offload history: %v\n", err) + } + } + + summary, err := c.summarizeHistory(ctx, toSummarize) + if err != nil { + return history, err + } + + compacted := make([]*ai.Message, 0, 2+len(history)-preserveFromIdx) + + if len(history) > 0 { + compacted = append(compacted, history[0]) + } + + // Build summary message with optional file reference + summaryMsg := fmt.Sprintf( + "[CONTEXT COMPACTION: The following is a summary of %d previous messages totaling ~%d tokens]\n\n%s\n\n", + len(toSummarize), + currentTokens-tokensToPreserve, + summary, + ) + + if historyFile != "" { + summaryMsg += fmt.Sprintf("[Full history preserved at: %s]\n\n", historyFile) + } + + summaryMsg += "[END SUMMARY - Continuing from recent context]" + + compacted = append(compacted, ai.NewUserTextMessage(summaryMsg)) + + compacted = append(compacted, history[preserveFromIdx:]...) + + return compacted, nil +} + +func (c *Compactor) summarizeHistory(ctx context.Context, messages []*ai.Message) (string, error) { + if c.genkitApp == nil { + return c.fallbackSummarize(messages), nil + } + + var historyText strings.Builder + for i, msg := range messages { + historyText.WriteString(fmt.Sprintf("\n--- Message %d (%s) ---\n", i+1, msg.Role)) + for _, part := range msg.Content { + if part.IsText() { + historyText.WriteString(part.Text) + historyText.WriteString("\n") + } else if part.IsToolRequest() { + historyText.WriteString(fmt.Sprintf("[Tool Call: %s]\n", part.ToolRequest.Name)) + } else if part.IsToolResponse() { + historyText.WriteString(fmt.Sprintf("[Tool Response: %s]\n", part.ToolResponse.Name)) + } + } + } + + prompt := fmt.Sprintf(`Summarize the following conversation history concisely, preserving: +1. Key decisions made +2. Important information discovered +3. Tool calls and their outcomes +4. Any errors encountered and their resolutions + +Be concise but comprehensive. Focus on information that would be needed to continue the task. + +CONVERSATION HISTORY: +%s + +SUMMARY:`, historyText.String()) + + resp, err := genkit.Generate(ctx, c.genkitApp, + ai.WithModelName(c.modelName), + ai.WithPrompt(prompt), + ) + if err != nil { + return c.fallbackSummarize(messages), nil + } + + return resp.Text(), nil +} + +func (c *Compactor) fallbackSummarize(messages []*ai.Message) string { + var summary strings.Builder + summary.WriteString("Previous conversation summary:\n") + + toolCalls := make(map[string]int) + for _, msg := range messages { + for _, part := range msg.Content { + if part.IsToolRequest() { + toolCalls[part.ToolRequest.Name]++ + } + } + } + + if len(toolCalls) > 0 { + summary.WriteString("\nTools used:\n") + for tool, count := range toolCalls { + summary.WriteString(fmt.Sprintf("- %s: %d times\n", tool, count)) + } + } + + summary.WriteString(fmt.Sprintf("\n[%d messages compacted]\n", len(messages))) + return summary.String() +} + +func (c *Compactor) ShouldCompact(history []*ai.Message) (bool, int, error) { + if !c.config.Enabled { + return false, 0, nil + } + + currentTokens, err := c.tokenCounter.CountTokens(history) + if err != nil { + return false, 0, err + } + + threshold := int(float64(c.contextWindow) * c.config.Threshold) + return currentTokens >= threshold, currentTokens, nil +} + +// offloadHistory saves conversation history to the backend before summarization +// This preserves the full history for debugging and retrieval +func (c *Compactor) offloadHistory(messages []*ai.Message) (string, error) { + if c.historyBackend == nil { + return "", nil + } + + // Create serializable structure for history + type HistoryEntry struct { + Timestamp time.Time `json:"timestamp"` + SessionID string `json:"session_id"` + Messages []struct { + Role string `json:"role"` + Content string `json:"content"` + Tools []struct { + Name string `json:"name"` + Type string `json:"type"` // "request" or "response" + Input any `json:"input,omitempty"` + Output any `json:"output,omitempty"` + } `json:"tools,omitempty"` + } `json:"messages"` + } + + entry := HistoryEntry{ + Timestamp: time.Now(), + SessionID: c.sessionID, + } + + for _, msg := range messages { + msgEntry := struct { + Role string `json:"role"` + Content string `json:"content"` + Tools []struct { + Name string `json:"name"` + Type string `json:"type"` + Input any `json:"input,omitempty"` + Output any `json:"output,omitempty"` + } `json:"tools,omitempty"` + }{ + Role: string(msg.Role), + } + + for _, part := range msg.Content { + if part.IsText() { + msgEntry.Content += part.Text + } else if part.IsToolRequest() { + msgEntry.Tools = append(msgEntry.Tools, struct { + Name string `json:"name"` + Type string `json:"type"` + Input any `json:"input,omitempty"` + Output any `json:"output,omitempty"` + }{ + Name: part.ToolRequest.Name, + Type: "request", + Input: part.ToolRequest.Input, + }) + } else if part.IsToolResponse() { + msgEntry.Tools = append(msgEntry.Tools, struct { + Name string `json:"name"` + Type string `json:"type"` + Input any `json:"input,omitempty"` + Output any `json:"output,omitempty"` + }{ + Name: part.ToolResponse.Name, + Type: "response", + Output: part.ToolResponse.Output, + }) + } + } + + entry.Messages = append(entry.Messages, msgEntry) + } + + data, err := json.MarshalIndent(entry, "", " ") + if err != nil { + return "", fmt.Errorf("failed to marshal history: %w", err) + } + + // Generate unique filename with timestamp + filename := fmt.Sprintf("%s/%s_%s.json", + c.historyPrefix, + c.sessionID, + time.Now().Format("20060102_150405"), + ) + + if err := c.historyBackend.WriteFile(filename, data); err != nil { + return "", fmt.Errorf("failed to write history: %w", err) + } + + return filename, nil +} + +// TruncateOldArgs truncates large tool arguments in older messages to reduce token count +// This is useful for tools like write_file and edit_file that have large inputs +func (c *Compactor) TruncateOldArgs(messages []*ai.Message, maxArgLength int) []*ai.Message { + if maxArgLength <= 0 { + maxArgLength = 500 // Default max length for truncated args + } + + // Don't modify the last few messages (keep them intact) + preserveCount := 4 + if len(messages) <= preserveCount { + return messages + } + + result := make([]*ai.Message, len(messages)) + copy(result[len(messages)-preserveCount:], messages[len(messages)-preserveCount:]) + + for i := 0; i < len(messages)-preserveCount; i++ { + msg := messages[i] + newContent := make([]*ai.Part, len(msg.Content)) + + for j, part := range msg.Content { + if part.IsToolRequest() { + // Check if this is a tool with potentially large args + toolName := part.ToolRequest.Name + if shouldTruncateArgs(toolName) { + // Create truncated version + truncatedInput := truncateInput(part.ToolRequest.Input, maxArgLength) + newPart := ai.NewToolRequestPart(&ai.ToolRequest{ + Name: toolName, + Input: truncatedInput, + }) + newContent[j] = newPart + } else { + newContent[j] = part + } + } else { + newContent[j] = part + } + } + + result[i] = &ai.Message{ + Role: msg.Role, + Content: newContent, + } + } + + return result +} + +// shouldTruncateArgs returns true for tools that typically have large arguments +func shouldTruncateArgs(toolName string) bool { + largeArgTools := map[string]bool{ + "write_file": true, + "edit_file": true, + "bash": true, + "create_file": true, + "update_file": true, + "patch_file": true, + "__write_file": true, + "__edit_file": true, + "__create_file": true, + } + return largeArgTools[toolName] +} + +// truncateInput truncates the input map values that are strings +func truncateInput(input any, maxLength int) any { + if input == nil { + return nil + } + + switch v := input.(type) { + case map[string]any: + result := make(map[string]any) + for key, val := range v { + if str, ok := val.(string); ok && len(str) > maxLength { + result[key] = str[:maxLength] + fmt.Sprintf("... [truncated %d chars]", len(str)-maxLength) + } else { + result[key] = val + } + } + return result + default: + return input + } +} diff --git a/pkg/harness/compaction_test.go b/pkg/harness/compaction_test.go new file mode 100644 index 00000000..5f6e4eed --- /dev/null +++ b/pkg/harness/compaction_test.go @@ -0,0 +1,295 @@ +package harness + +import ( + "context" + "strings" + "testing" + + "github.com/firebase/genkit/go/ai" +) + +// mockHistoryBackend implements HistoryBackend for testing +type mockHistoryBackend struct { + files map[string][]byte +} + +func newMockHistoryBackend() *mockHistoryBackend { + return &mockHistoryBackend{ + files: make(map[string][]byte), + } +} + +func (m *mockHistoryBackend) WriteFile(path string, content []byte) error { + m.files[path] = content + return nil +} + +func (m *mockHistoryBackend) ReadFile(path string) ([]byte, error) { + return m.files[path], nil +} + +func TestCompactor_TruncateOldArgs(t *testing.T) { + compactor := NewCompactor(nil, "", CompactionConfig{Enabled: true}, 100000) + + // Create messages with large tool arguments + messages := []*ai.Message{ + ai.NewUserTextMessage("Hello"), + { + Role: ai.RoleModel, + Content: []*ai.Part{ + ai.NewToolRequestPart(&ai.ToolRequest{ + Name: "write_file", + Input: map[string]any{ + "path": "/some/path.txt", + "content": "This is a very long content string that should be truncated when it exceeds the maximum allowed length for old tool arguments. " + string(make([]byte, 1000)), + }, + }), + }, + }, + ai.NewUserTextMessage("Continue"), + ai.NewModelTextMessage("Done"), + ai.NewUserTextMessage("More"), + ai.NewModelTextMessage("Finished"), + } + + // Truncate with max length of 100 + truncated := compactor.TruncateOldArgs(messages, 100) + + // Last 4 messages should be unchanged + if len(truncated) != len(messages) { + t.Errorf("Expected %d messages, got %d", len(messages), len(truncated)) + } + + // Check that the old write_file tool call was truncated + if len(truncated[1].Content) > 0 && truncated[1].Content[0].IsToolRequest() { + input := truncated[1].Content[0].ToolRequest.Input.(map[string]any) + content := input["content"].(string) + if len(content) > 200 { // 100 chars + truncation message + t.Errorf("Content should be truncated, got length %d", len(content)) + } + } +} + +func TestCompactor_ShouldTruncateArgs(t *testing.T) { + tests := []struct { + toolName string + expected bool + }{ + {"write_file", true}, + {"edit_file", true}, + {"bash", true}, + {"read_file", false}, + {"list_directory", false}, + {"__write_file", true}, + {"unknown_tool", false}, + } + + for _, tt := range tests { + t.Run(tt.toolName, func(t *testing.T) { + result := shouldTruncateArgs(tt.toolName) + if result != tt.expected { + t.Errorf("shouldTruncateArgs(%q) = %v, want %v", tt.toolName, result, tt.expected) + } + }) + } +} + +func TestCompactor_TruncateInput(t *testing.T) { + t.Run("truncate long string", func(t *testing.T) { + input := map[string]any{ + "content": "This is a very long string that will be truncated", + } + result := truncateInput(input, 10) + resultMap := result.(map[string]any) + content := resultMap["content"].(string) + + // Should start with first 10 chars and contain truncation message + if !strings.Contains(content, "This is a ") { + t.Errorf("Truncated content should start with original prefix, got: %s", content) + } + if !strings.Contains(content, "[truncated") { + t.Errorf("Truncated content should contain truncation message, got: %s", content) + } + }) + + t.Run("short string unchanged", func(t *testing.T) { + input := map[string]any{ + "content": "short", + } + result := truncateInput(input, 100) + resultMap := result.(map[string]any) + content := resultMap["content"].(string) + + if content != "short" { + t.Errorf("Short content should be unchanged, got: %s", content) + } + }) + + t.Run("nil input", func(t *testing.T) { + result := truncateInput(nil, 100) + if result != nil { + t.Error("Expected nil result for nil input") + } + }) +} + +func TestCompactor_WithHistoryBackend(t *testing.T) { + backend := newMockHistoryBackend() + compactor := NewCompactor(nil, "", CompactionConfig{ + Enabled: true, + HistoryOffload: true, + }, 100000) + + compactor.WithHistoryBackend(backend, "/history", "session-123") + + if compactor.historyBackend != backend { + t.Error("historyBackend not set correctly") + } + if compactor.historyPrefix != "/history" { + t.Errorf("historyPrefix = %q, want /history", compactor.historyPrefix) + } + if compactor.sessionID != "session-123" { + t.Errorf("sessionID = %q, want session-123", compactor.sessionID) + } +} + +func TestCompactor_OffloadHistory(t *testing.T) { + backend := newMockHistoryBackend() + compactor := NewCompactor(nil, "", CompactionConfig{ + Enabled: true, + HistoryOffload: true, + }, 100000) + compactor.WithHistoryBackend(backend, "/history", "test-session") + + messages := []*ai.Message{ + ai.NewUserTextMessage("Hello world"), + ai.NewModelTextMessage("Hi there!"), + { + Role: ai.RoleModel, + Content: []*ai.Part{ + ai.NewToolRequestPart(&ai.ToolRequest{ + Name: "read_file", + Input: map[string]any{"path": "/test.txt"}, + }), + }, + }, + } + + filename, err := compactor.offloadHistory(messages) + if err != nil { + t.Fatalf("offloadHistory() error = %v", err) + } + + if filename == "" { + t.Error("Expected non-empty filename") + } + + // Check that file was written + if len(backend.files) == 0 { + t.Error("No files written to backend") + } + + // Verify file content contains expected data + for _, content := range backend.files { + if len(content) == 0 { + t.Error("Empty file content") + } + // Basic JSON validation + if content[0] != '{' { + t.Error("Expected JSON content starting with {") + } + } +} + +func TestCompactor_OffloadHistory_NoBackend(t *testing.T) { + compactor := NewCompactor(nil, "", CompactionConfig{ + Enabled: true, + HistoryOffload: true, + }, 100000) + + messages := []*ai.Message{ + ai.NewUserTextMessage("Hello"), + } + + filename, err := compactor.offloadHistory(messages) + if err != nil { + t.Errorf("offloadHistory() error = %v", err) + } + if filename != "" { + t.Errorf("Expected empty filename when no backend, got %q", filename) + } +} + +func TestCompactor_CompactIfNeeded_Disabled(t *testing.T) { + compactor := NewCompactor(nil, "", CompactionConfig{Enabled: false}, 100000) + + messages := []*ai.Message{ + ai.NewUserTextMessage("Hello"), + } + + result, compacted, err := compactor.CompactIfNeeded(context.Background(), messages) + if err != nil { + t.Fatalf("CompactIfNeeded() error = %v", err) + } + if compacted { + t.Error("Should not compact when disabled") + } + if len(result) != len(messages) { + t.Error("Messages should be unchanged when disabled") + } +} + +func TestCompactor_ShouldCompact(t *testing.T) { + compactor := NewCompactor(nil, "", CompactionConfig{ + Enabled: true, + Threshold: 0.5, + }, 1000) // Small context window for testing + + tests := []struct { + name string + msgCount int + wantCompact bool + }{ + {"small history", 2, false}, + {"large history", 100, true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + messages := make([]*ai.Message, tt.msgCount) + for i := 0; i < tt.msgCount; i++ { + messages[i] = ai.NewUserTextMessage("This is message number " + string(rune('A'+i%26)) + " with some content") + } + + should, _, err := compactor.ShouldCompact(messages) + if err != nil { + t.Fatalf("ShouldCompact() error = %v", err) + } + if should != tt.wantCompact { + t.Errorf("ShouldCompact() = %v, want %v", should, tt.wantCompact) + } + }) + } +} + +func TestSimpleTokenCounter(t *testing.T) { + counter := &SimpleTokenCounter{} + + messages := []*ai.Message{ + ai.NewUserTextMessage("Hello world this is a test message"), + ai.NewModelTextMessage("This is the response with more words"), + } + + count, err := counter.CountTokens(messages) + if err != nil { + t.Fatalf("CountTokens() error = %v", err) + } + + // Simple token counter: words * 4 / 3 + // "Hello world this is a test message" = 7 words + // "This is the response with more words" = 7 words + // 14 words * 4 / 3 = 18 tokens (approx) + if count < 10 || count > 30 { + t.Errorf("CountTokens() = %d, expected between 10-30", count) + } +} diff --git a/pkg/harness/config.go b/pkg/harness/config.go new file mode 100644 index 00000000..3e90f6dc --- /dev/null +++ b/pkg/harness/config.go @@ -0,0 +1,297 @@ +// Package harness provides an agentic execution harness for Station. +// It implements a Claude Agent SDK-like execution model with manual loop control, +// pre/post tool hooks, doom loop detection, and context compaction. +package harness + +import ( + "time" +) + +// HarnessConfig is the top-level configuration for the agentic harness. +// It is loaded from stn.yaml under the "harness" key. +type HarnessConfig struct { + Workspace WorkspaceConfig `yaml:"workspace" json:"workspace"` + Compaction CompactionConfig `yaml:"compaction" json:"compaction"` + Git GitConfig `yaml:"git" json:"git"` + NATS NATSConfig `yaml:"nats" json:"nats"` + Permissions PermissionsConfig `yaml:"permissions" json:"permissions"` + Streaming StreamingConfig `yaml:"streaming" json:"streaming"` +} + +// WorkspaceConfig defines workspace settings. +type WorkspaceConfig struct { + // Path is the workspace directory path (relative to station root). + // Default: "./workspace" + Path string `yaml:"path" json:"path"` + + // Mode is the workspace isolation mode. + // Options: "host" (raw filesystem), "sandbox" (Docker), "agentfs" (SQLite) + // Default: "host" + Mode string `yaml:"mode" json:"mode"` + + // IsolationMode controls how workspaces are isolated for concurrent runs. + // Options: + // - "shared": All runs share the base workspace (NOT RECOMMENDED for concurrent) + // - "per_run": Each agent run gets its own subdirectory (basePath//) + // - "per_workflow": Workflow steps share workspace, isolated between workflows + // Default: "per_workflow" + IsolationMode string `yaml:"isolation_mode" json:"isolation_mode"` + + // CleanupOnComplete removes the workspace after successful completion. + // Default: false + CleanupOnComplete bool `yaml:"cleanup_on_complete" json:"cleanup_on_complete"` + + // CleanupAfter specifies duration after which to cleanup completed workspaces. + // Only applies when CleanupOnComplete is false. + // Default: "24h" + CleanupAfter string `yaml:"cleanup_after" json:"cleanup_after"` +} + +// GitSourceConfig defines settings for cloning a git repo into the workspace. +// This allows agents to work on specific repositories. +type GitSourceConfig struct { + // URL is the git repository URL to clone. + // Supports HTTPS and SSH URLs. + // Example: "https://github.com/org/repo.git" or "git@github.com:org/repo.git" + URL string `yaml:"url" json:"url"` + + // Branch is the branch to checkout after cloning. + // Default: repository's default branch (usually "main") + Branch string `yaml:"branch" json:"branch"` + + // Ref is a specific commit SHA, tag, or ref to checkout. + // Takes precedence over Branch if both are specified. + Ref string `yaml:"ref" json:"ref"` + + // Depth for shallow clone (0 = full clone). + // Default: 1 (shallow clone for faster setup) + Depth int `yaml:"depth" json:"depth"` + + // Submodules controls whether to initialize submodules. + // Default: false + Submodules bool `yaml:"submodules" json:"submodules"` + + // SSHKeyPath is the path to SSH key for private repos. + // If empty, uses default SSH agent. + SSHKeyPath string `yaml:"ssh_key_path" json:"ssh_key_path"` +} + +// ExecutionContext provides runtime context for workspace isolation. +// This is passed at execution time, not from config. +type ExecutionContext struct { + // SessionID is a user-provided identifier for workspace persistence across runs. + // If provided, takes precedence over WorkflowRunID and AgentRunID. + // Use this to continue work from a previous run. + SessionID string + // WorkflowRunID is the workflow run ID (for per_workflow isolation). + WorkflowRunID string + // AgentRunID is the agent run ID (for per_run isolation). + AgentRunID string + // AgentName for branch naming. + AgentName string + // Task description for branch naming. + Task string +} + +// CompactionConfig defines context compaction settings. +type CompactionConfig struct { + // Enabled enables automatic context compaction. + // Default: true + Enabled bool `yaml:"enabled" json:"enabled"` + + // Threshold is the percentage of context window that triggers compaction. + // Default: 0.85 (85%) + Threshold float64 `yaml:"threshold" json:"threshold"` + + // ProtectTokens is the number of recent tokens to keep protected from pruning. + // Default: 40000 + ProtectTokens int `yaml:"protect_tokens" json:"protect_tokens"` + + // HistoryOffload enables saving full conversation history before summarization. + // This preserves the complete history for debugging and retrieval. + // Default: false + HistoryOffload bool `yaml:"history_offload" json:"history_offload"` +} + +// GitConfig defines git workflow settings. +type GitConfig struct { + // AutoBranch enables automatic branch creation for each task. + // Default: true + AutoBranch bool `yaml:"auto_branch" json:"auto_branch"` + + // BranchPrefix is the prefix for agent-created branches. + // Default: "agent/" + BranchPrefix string `yaml:"branch_prefix" json:"branch_prefix"` + + // AutoCommit enables automatic commit on task completion. + // Default: false + AutoCommit bool `yaml:"auto_commit" json:"auto_commit"` + + // RequireApproval requires human approval before push/PR. + // Default: true + RequireApproval bool `yaml:"require_approval" json:"require_approval"` + + // WorkflowBranchStrategy defines how branches are managed in workflows. + // Options: "shared" (all agents on same branch), "per_agent" (new branch per agent) + // Default: "shared" + WorkflowBranchStrategy string `yaml:"workflow_branch_strategy" json:"workflow_branch_strategy"` +} + +// NATSConfig defines NATS integration settings for inter-agent communication. +type NATSConfig struct { + // Enabled enables NATS integration. + // Default: true + Enabled bool `yaml:"enabled" json:"enabled"` + + // KVBucket is the NATS KV bucket name for state. + // Default: "harness-state" + KVBucket string `yaml:"kv_bucket" json:"kv_bucket"` + + // ObjectBucket is the NATS Object Store bucket name for files. + // Default: "harness-files" + ObjectBucket string `yaml:"object_bucket" json:"object_bucket"` + + // MaxFileSize is the maximum file size for object store. + // Default: "100MB" + MaxFileSize string `yaml:"max_file_size" json:"max_file_size"` + + // TTL is the time-to-live for stored artifacts. + // Default: "24h" + TTL string `yaml:"ttl" json:"ttl"` +} + +// PermissionsConfig defines command permissions (OpenCode-style). +type PermissionsConfig struct { + // ExternalDirectory controls access to paths outside workspace. + // Options: "deny", "ask", "allow" + // Default: "deny" + ExternalDirectory string `yaml:"external_directory" json:"external_directory"` + + // Bash is a map of command patterns to actions. + // Actions: "allow", "deny", "ask" + // Patterns support wildcards: "*" matches any sequence + // Example: {"*": "allow", "rm -rf *": "deny", "git push *": "ask"} + Bash map[string]string `yaml:"bash" json:"bash"` +} + +type StreamingConfig struct { + Enabled bool `yaml:"enabled" json:"enabled"` +} + +// AgentHarnessConfig is per-agent configuration that overrides global settings. +type AgentHarnessConfig struct { + // MaxSteps is the maximum number of agentic loop iterations. + // Default: 50 + MaxSteps int `yaml:"max_steps" json:"max_steps"` + + // DoomLoopThreshold is the number of identical consecutive tool calls + // before triggering doom loop detection. + // Default: 3 + DoomLoopThreshold int `yaml:"doom_loop_threshold" json:"doom_loop_threshold"` + + // Timeout is the maximum execution time for the entire task. + // Default: 30m + Timeout time.Duration `yaml:"timeout" json:"timeout"` +} + +// DefaultHarnessConfig returns the default harness configuration. +func DefaultHarnessConfig() *HarnessConfig { + return &HarnessConfig{ + Workspace: WorkspaceConfig{ + Path: "./workspace", + Mode: "host", + IsolationMode: "per_workflow", + CleanupOnComplete: false, + CleanupAfter: "24h", + }, + Compaction: CompactionConfig{ + Enabled: true, + Threshold: 0.85, + ProtectTokens: 40000, + }, + Git: GitConfig{ + AutoBranch: true, + BranchPrefix: "agent/", + AutoCommit: false, + RequireApproval: true, + WorkflowBranchStrategy: "shared", + }, + NATS: NATSConfig{ + Enabled: true, + KVBucket: "harness-state", + ObjectBucket: "harness-files", + MaxFileSize: "100MB", + TTL: "24h", + }, + Permissions: PermissionsConfig{ + ExternalDirectory: "deny", + Bash: map[string]string{ + "*": "allow", + "rm -rf *": "deny", + "rm -r *": "deny", + "rm -f /*": "deny", + "chmod 777 *": "deny", + "chmod -R 777 *": "deny", + "curl * | sh": "deny", + "curl * | bash": "deny", + "wget * | sh": "deny", + "wget * | bash": "deny", + "find * -delete*": "deny", + "find * -exec rm *": "deny", + "git push --force*": "deny", + "git reset --hard*": "deny", + "git push *": "ask", + "git commit *": "ask", + }, + }, + } +} + +// DefaultAgentHarnessConfig returns the default per-agent configuration. +func DefaultAgentHarnessConfig() *AgentHarnessConfig { + return &AgentHarnessConfig{ + MaxSteps: 50, + DoomLoopThreshold: 3, + Timeout: 30 * time.Minute, + } +} + +// Merge merges agent-specific config with global config. +func (c *AgentHarnessConfig) Merge(global *HarnessConfig) *AgentHarnessConfig { + merged := *c + if merged.MaxSteps == 0 { + merged.MaxSteps = DefaultAgentHarnessConfig().MaxSteps + } + if merged.DoomLoopThreshold == 0 { + merged.DoomLoopThreshold = DefaultAgentHarnessConfig().DoomLoopThreshold + } + if merged.Timeout == 0 { + merged.Timeout = DefaultAgentHarnessConfig().Timeout + } + return &merged +} + +// ResolveWorkspacePath determines the workspace path based on isolation mode and execution context. +func ResolveWorkspacePath(basePath string, isolationMode string, execCtx *ExecutionContext) string { + if execCtx == nil { + return basePath + } + + switch isolationMode { + case "per_workflow": + if execCtx.WorkflowRunID != "" { + return basePath + "/workflow/" + execCtx.WorkflowRunID + } + if execCtx.AgentRunID != "" { + return basePath + "/agent/" + execCtx.AgentRunID + } + case "per_run": + if execCtx.AgentRunID != "" { + return basePath + "/run/" + execCtx.AgentRunID + } + case "shared": + return basePath + } + + return basePath +} diff --git a/pkg/harness/config_test.go b/pkg/harness/config_test.go new file mode 100644 index 00000000..ed5ab13d --- /dev/null +++ b/pkg/harness/config_test.go @@ -0,0 +1,107 @@ +package harness + +import ( + "testing" + "time" +) + +func TestDefaultHarnessConfig(t *testing.T) { + cfg := DefaultHarnessConfig() + + if cfg.Workspace.Path != "./workspace" { + t.Errorf("Workspace.Path = %q, want %q", cfg.Workspace.Path, "./workspace") + } + if cfg.Workspace.Mode != "host" { + t.Errorf("Workspace.Mode = %q, want %q", cfg.Workspace.Mode, "host") + } + + if !cfg.Compaction.Enabled { + t.Error("Compaction.Enabled should be true by default") + } + if cfg.Compaction.Threshold != 0.85 { + t.Errorf("Compaction.Threshold = %f, want %f", cfg.Compaction.Threshold, 0.85) + } + if cfg.Compaction.ProtectTokens != 40000 { + t.Errorf("Compaction.ProtectTokens = %d, want %d", cfg.Compaction.ProtectTokens, 40000) + } + + if !cfg.Git.AutoBranch { + t.Error("Git.AutoBranch should be true by default") + } + if cfg.Git.BranchPrefix != "agent/" { + t.Errorf("Git.BranchPrefix = %q, want %q", cfg.Git.BranchPrefix, "agent/") + } + if cfg.Git.WorkflowBranchStrategy != "shared" { + t.Errorf("Git.WorkflowBranchStrategy = %q, want %q", cfg.Git.WorkflowBranchStrategy, "shared") + } + + if !cfg.NATS.Enabled { + t.Error("NATS.Enabled should be true by default") + } + if cfg.NATS.KVBucket != "harness-state" { + t.Errorf("NATS.KVBucket = %q, want %q", cfg.NATS.KVBucket, "harness-state") + } + + if cfg.Permissions.ExternalDirectory != "deny" { + t.Errorf("Permissions.ExternalDirectory = %q, want %q", cfg.Permissions.ExternalDirectory, "deny") + } + + if action, ok := cfg.Permissions.Bash["rm -rf *"]; !ok || action != "deny" { + t.Error("rm -rf * should be denied by default") + } + if action, ok := cfg.Permissions.Bash["git push *"]; !ok || action != "ask" { + t.Error("git push * should require ask by default") + } +} + +func TestDefaultAgentHarnessConfig(t *testing.T) { + cfg := DefaultAgentHarnessConfig() + + if cfg.MaxSteps != 50 { + t.Errorf("MaxSteps = %d, want %d", cfg.MaxSteps, 50) + } + if cfg.DoomLoopThreshold != 3 { + t.Errorf("DoomLoopThreshold = %d, want %d", cfg.DoomLoopThreshold, 3) + } + if cfg.Timeout != 30*time.Minute { + t.Errorf("Timeout = %v, want %v", cfg.Timeout, 30*time.Minute) + } +} + +func TestAgentHarnessConfig_Merge(t *testing.T) { + global := DefaultHarnessConfig() + + agentCfg := &AgentHarnessConfig{ + MaxSteps: 100, + } + + merged := agentCfg.Merge(global) + + if merged.MaxSteps != 100 { + t.Errorf("merged.MaxSteps = %d, want %d", merged.MaxSteps, 100) + } + if merged.DoomLoopThreshold != 3 { + t.Errorf("merged.DoomLoopThreshold = %d, want %d (default)", merged.DoomLoopThreshold, 3) + } + if merged.Timeout != 30*time.Minute { + t.Errorf("merged.Timeout = %v, want %v (default)", merged.Timeout, 30*time.Minute) + } +} + +func TestAgentHarnessConfig_MergeWithZeroValues(t *testing.T) { + global := DefaultHarnessConfig() + + agentCfg := &AgentHarnessConfig{} + + merged := agentCfg.Merge(global) + + if merged.MaxSteps != 50 { + t.Errorf("merged.MaxSteps = %d, want %d (should use default)", merged.MaxSteps, 50) + } + if merged.DoomLoopThreshold != 3 { + t.Errorf("merged.DoomLoopThreshold = %d, want %d (should use default)", merged.DoomLoopThreshold, 3) + } + if merged.Timeout != 30*time.Minute { + t.Errorf("merged.Timeout = %v, want %v (should use default)", merged.Timeout, 30*time.Minute) + } +} diff --git a/pkg/harness/doom_loop.go b/pkg/harness/doom_loop.go new file mode 100644 index 00000000..0ae48e1a --- /dev/null +++ b/pkg/harness/doom_loop.go @@ -0,0 +1,100 @@ +package harness + +import ( + "crypto/sha256" + "encoding/hex" + "encoding/json" + "sync" +) + +type DoomLoopDetector struct { + mu sync.Mutex + threshold int + callHistory []callRecord + consecutiveRuns map[string]int +} + +type callRecord struct { + toolName string + argHash string +} + +func NewDoomLoopDetector(threshold int) *DoomLoopDetector { + if threshold <= 0 { + threshold = 3 + } + return &DoomLoopDetector{ + threshold: threshold, + callHistory: make([]callRecord, 0), + consecutiveRuns: make(map[string]int), + } +} + +func (d *DoomLoopDetector) IsInDoomLoop(toolName string, args interface{}) bool { + d.mu.Lock() + defer d.mu.Unlock() + + argHash := hashArgs(args) + key := toolName + ":" + argHash + + if len(d.callHistory) > 0 { + last := d.callHistory[len(d.callHistory)-1] + lastKey := last.toolName + ":" + last.argHash + + if lastKey == key { + d.consecutiveRuns[key]++ + } else { + d.consecutiveRuns[key] = 1 + } + } else { + d.consecutiveRuns[key] = 1 + } + + return d.consecutiveRuns[key] >= d.threshold +} + +func (d *DoomLoopDetector) Record(toolName string, args interface{}) { + d.mu.Lock() + defer d.mu.Unlock() + + argHash := hashArgs(args) + d.callHistory = append(d.callHistory, callRecord{ + toolName: toolName, + argHash: argHash, + }) + + const maxHistory = 100 + if len(d.callHistory) > maxHistory { + d.callHistory = d.callHistory[len(d.callHistory)-maxHistory:] + } +} + +func (d *DoomLoopDetector) Reset() { + d.mu.Lock() + defer d.mu.Unlock() + d.callHistory = make([]callRecord, 0) + d.consecutiveRuns = make(map[string]int) +} + +func (d *DoomLoopDetector) GetConsecutiveCount(toolName string, args interface{}) int { + d.mu.Lock() + defer d.mu.Unlock() + + argHash := hashArgs(args) + key := toolName + ":" + argHash + return d.consecutiveRuns[key] +} + +func hashArgs(args interface{}) string { + if args == nil { + return "nil" + } + + data, err := json.Marshal(args) + if err != nil { + return "error" + } + + hash := sha256.Sum256(data) + return hex.EncodeToString(hash[:8]) +} diff --git a/pkg/harness/doom_loop_test.go b/pkg/harness/doom_loop_test.go new file mode 100644 index 00000000..2fc4132d --- /dev/null +++ b/pkg/harness/doom_loop_test.go @@ -0,0 +1,144 @@ +package harness + +import ( + "testing" +) + +func TestDoomLoopDetector_NewDetector(t *testing.T) { + detector := NewDoomLoopDetector(3) + if detector.threshold != 3 { + t.Errorf("threshold = %d, want %d", detector.threshold, 3) + } + + detectorDefault := NewDoomLoopDetector(0) + if detectorDefault.threshold != 3 { + t.Errorf("default threshold = %d, want %d", detectorDefault.threshold, 3) + } +} + +func TestDoomLoopDetector_DetectsLoop(t *testing.T) { + detector := NewDoomLoopDetector(3) + + args := map[string]interface{}{"file": "test.txt"} + + if detector.IsInDoomLoop("read_file", args) { + t.Error("Should not detect loop on first call") + } + detector.Record("read_file", args) + + if detector.IsInDoomLoop("read_file", args) { + t.Error("Should not detect loop on second call") + } + detector.Record("read_file", args) + + if !detector.IsInDoomLoop("read_file", args) { + t.Error("Should detect loop on third consecutive call") + } +} + +func TestDoomLoopDetector_DifferentArgs(t *testing.T) { + detector := NewDoomLoopDetector(3) + + args1 := map[string]interface{}{"file": "test1.txt"} + args2 := map[string]interface{}{"file": "test2.txt"} + + detector.IsInDoomLoop("read_file", args1) + detector.Record("read_file", args1) + + detector.IsInDoomLoop("read_file", args2) + detector.Record("read_file", args2) + + detector.IsInDoomLoop("read_file", args1) + detector.Record("read_file", args1) + + if detector.IsInDoomLoop("read_file", args2) { + t.Error("Should not detect loop with different args interleaved") + } +} + +func TestDoomLoopDetector_DifferentTools(t *testing.T) { + detector := NewDoomLoopDetector(3) + + args := map[string]interface{}{"path": "/"} + + detector.IsInDoomLoop("list_files", args) + detector.Record("list_files", args) + + detector.IsInDoomLoop("read_file", args) + detector.Record("read_file", args) + + detector.IsInDoomLoop("list_files", args) + detector.Record("list_files", args) + + if detector.IsInDoomLoop("read_file", args) { + t.Error("Should not detect loop with different tools interleaved") + } +} + +func TestDoomLoopDetector_Reset(t *testing.T) { + detector := NewDoomLoopDetector(3) + + args := map[string]interface{}{"file": "test.txt"} + + detector.IsInDoomLoop("read_file", args) + detector.Record("read_file", args) + detector.IsInDoomLoop("read_file", args) + detector.Record("read_file", args) + + detector.Reset() + + if detector.IsInDoomLoop("read_file", args) { + t.Error("Should not detect loop after reset") + } +} + +func TestDoomLoopDetector_GetConsecutiveCount(t *testing.T) { + detector := NewDoomLoopDetector(5) + + args := map[string]interface{}{"cmd": "ls"} + + detector.IsInDoomLoop("bash", args) + detector.Record("bash", args) + detector.IsInDoomLoop("bash", args) + detector.Record("bash", args) + detector.IsInDoomLoop("bash", args) + + count := detector.GetConsecutiveCount("bash", args) + if count != 3 { + t.Errorf("consecutive count = %d, want %d", count, 3) + } +} + +func TestHashArgs(t *testing.T) { + tests := []struct { + name string + args interface{} + stable bool + }{ + {"nil", nil, true}, + {"empty map", map[string]interface{}{}, true}, + {"simple map", map[string]interface{}{"key": "value"}, true}, + {"nested map", map[string]interface{}{"outer": map[string]interface{}{"inner": "value"}}, true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + hash1 := hashArgs(tt.args) + hash2 := hashArgs(tt.args) + + if tt.stable && hash1 != hash2 { + t.Errorf("hash not stable for %v: %s != %s", tt.args, hash1, hash2) + } + + if hash1 == "" { + t.Error("hash should not be empty") + } + }) + } + + hash1 := hashArgs(map[string]interface{}{"a": 1}) + hash2 := hashArgs(map[string]interface{}{"b": 2}) + if hash1 == hash2 { + t.Error("different args should have different hashes") + } +} diff --git a/pkg/harness/executor.go b/pkg/harness/executor.go new file mode 100644 index 00000000..9fb3606f --- /dev/null +++ b/pkg/harness/executor.go @@ -0,0 +1,698 @@ +package harness + +import ( + "context" + "fmt" + "strings" + "time" + + "station/pkg/harness/memory" + "station/pkg/harness/sandbox" + "station/pkg/harness/skills" + "station/pkg/harness/stream" + + "github.com/firebase/genkit/go/ai" + "github.com/firebase/genkit/go/genkit" + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/trace" +) + +var tracer = otel.Tracer("station.harness") + +type ExecutionResult struct { + Success bool `json:"success"` + Response string `json:"response"` + Error string `json:"error,omitempty"` + TotalSteps int `json:"total_steps"` + TotalTokens int `json:"total_tokens"` + FinishReason string `json:"finish_reason"` + Duration time.Duration `json:"duration"` + Metadata map[string]interface{} `json:"metadata,omitempty"` + History []*ai.Message `json:"-"` // Full message history for session persistence +} + +type AgenticExecutor struct { + genkitApp *genkit.Genkit + config *HarnessConfig + agentConfig *AgentHarnessConfig + hooks *HookRegistry + doomDetector *DoomLoopDetector + compactor *Compactor + workspace WorkspaceManager + gitManager GitManager + promptBuilder *PromptBuilder + modelName string + systemPrompt string + streamCtx *stream.StreamContext + publisher stream.Publisher + sandbox sandbox.Sandbox + sandboxConfig *sandbox.Config + skillsMiddleware *skills.SkillsMiddleware + memoryMiddleware *memory.MemoryMiddleware + initialHistory []*ai.Message +} + +type ExecutorOption func(*AgenticExecutor) + +func WithHooks(hooks *HookRegistry) ExecutorOption { + return func(e *AgenticExecutor) { + e.hooks = hooks + } +} + +func WithCompactor(compactor *Compactor) ExecutorOption { + return func(e *AgenticExecutor) { + e.compactor = compactor + } +} + +func WithWorkspace(workspace WorkspaceManager) ExecutorOption { + return func(e *AgenticExecutor) { + e.workspace = workspace + } +} + +func WithGitManager(gitManager GitManager) ExecutorOption { + return func(e *AgenticExecutor) { + e.gitManager = gitManager + } +} + +func WithModelName(modelName string) ExecutorOption { + return func(e *AgenticExecutor) { + e.modelName = modelName + } +} + +func WithSystemPrompt(systemPrompt string) ExecutorOption { + return func(e *AgenticExecutor) { + e.systemPrompt = systemPrompt + } +} + +func WithStreamPublisher(publisher stream.Publisher) ExecutorOption { + return func(e *AgenticExecutor) { + e.publisher = publisher + } +} + +func WithSandbox(sb sandbox.Sandbox) ExecutorOption { + return func(e *AgenticExecutor) { + e.sandbox = sb + } +} + +func WithSandboxConfig(cfg *sandbox.Config) ExecutorOption { + return func(e *AgenticExecutor) { + e.sandboxConfig = cfg + } +} + +func WithSkillsMiddleware(sm *skills.SkillsMiddleware) ExecutorOption { + return func(e *AgenticExecutor) { + e.skillsMiddleware = sm + } +} + +func WithMemoryMiddleware(mm *memory.MemoryMiddleware) ExecutorOption { + return func(e *AgenticExecutor) { + e.memoryMiddleware = mm + } +} + +// WithInitialHistory sets the initial message history for session persistence. +// This allows REPL-style conversations where the agent can continue from +// previous interactions with full context. +func WithInitialHistory(history []*ai.Message) ExecutorOption { + return func(e *AgenticExecutor) { + e.initialHistory = history + } +} + +func NewAgenticExecutor( + genkitApp *genkit.Genkit, + config *HarnessConfig, + agentConfig *AgentHarnessConfig, + opts ...ExecutorOption, +) *AgenticExecutor { + if config == nil { + config = DefaultHarnessConfig() + } + if agentConfig == nil { + agentConfig = DefaultAgentHarnessConfig() + } + + e := &AgenticExecutor{ + genkitApp: genkitApp, + config: config, + agentConfig: agentConfig.Merge(config), + hooks: NewHookRegistry(), + doomDetector: NewDoomLoopDetector(agentConfig.DoomLoopThreshold), + modelName: "openai/gpt-4o-mini", + } + + e.hooks.RegisterDoomLoopDetection(e.doomDetector) + e.hooks.RegisterPermissionCheck(config.Permissions) + + for _, opt := range opts { + opt(e) + } + + return e +} + +type ExecuteOptions struct { + StationRunID string + RunUUID string + WorkflowRunID string + SessionID string + AgentName string + StationID string +} + +func (e *AgenticExecutor) Execute( + ctx context.Context, + agentID string, + task string, + tools []ai.Tool, +) (*ExecutionResult, error) { + return e.ExecuteWithOptions(ctx, agentID, task, tools, ExecuteOptions{}) +} + +func (e *AgenticExecutor) ExecuteWithOptions( + ctx context.Context, + agentID string, + task string, + tools []ai.Tool, + opts ExecuteOptions, +) (*ExecutionResult, error) { + ctx, span := tracer.Start(ctx, "agent_execution", + trace.WithAttributes( + attribute.String("agent_id", agentID), + attribute.String("harness", "agentic"), + attribute.Int("max_steps", e.agentConfig.MaxSteps), + ), + ) + defer span.End() + + startTime := time.Now() + + e.initializeStreaming(ctx, agentID, task, opts, span) + + setupCtx, setupSpan := tracer.Start(ctx, "harness_setup") + if err := e.setup(setupCtx, agentID, task); err != nil { + setupSpan.RecordError(err) + setupSpan.End() + result := &ExecutionResult{ + Success: false, + Error: fmt.Sprintf("setup failed: %v", err), + FinishReason: "setup_error", + Duration: time.Since(startTime), + } + e.emitRunComplete(ctx, result, err) + return result, err + } + setupSpan.End() + + loopCtx, loopSpan := tracer.Start(ctx, "agentic_loop") + result, err := e.runLoop(loopCtx, task, tools) + if err != nil { + loopSpan.RecordError(err) + } + loopSpan.SetAttributes( + attribute.Int("total_steps", result.TotalSteps), + attribute.Int("total_tokens", result.TotalTokens), + attribute.String("finish_reason", result.FinishReason), + ) + loopSpan.End() + + cleanupCtx, cleanupSpan := tracer.Start(ctx, "harness_cleanup") + if cleanupErr := e.cleanup(cleanupCtx, agentID, result); cleanupErr != nil { + cleanupSpan.RecordError(cleanupErr) + } + cleanupSpan.End() + + result.Duration = time.Since(startTime) + span.SetAttributes( + attribute.Bool("success", result.Success), + attribute.Int64("duration_ms", result.Duration.Milliseconds()), + ) + + e.emitRunComplete(ctx, result, err) + + return result, err +} + +func (e *AgenticExecutor) initializeStreaming(ctx context.Context, agentID, task string, opts ExecuteOptions, span trace.Span) { + if e.publisher == nil { + return + } + + runUUID := opts.RunUUID + if runUUID == "" { + runUUID = fmt.Sprintf("%s-%d", agentID, time.Now().UnixNano()) + } + + agentName := opts.AgentName + if agentName == "" { + agentName = agentID + } + + ids := stream.StreamIdentifiers{ + StationRunID: opts.StationRunID, + RunUUID: runUUID, + WorkflowRunID: opts.WorkflowRunID, + SessionID: opts.SessionID, + AgentID: agentID, + AgentName: agentName, + StationID: opts.StationID, + } + + e.streamCtx = stream.NewStreamContext(ids, e.publisher) + + if err := e.streamCtx.EmitRunStart(ctx, agentID, agentName, task, e.agentConfig.MaxSteps); err != nil { + span.RecordError(err) + } +} + +func (e *AgenticExecutor) emitRunComplete(ctx context.Context, result *ExecutionResult, err error) { + if e.streamCtx != nil { + e.streamCtx.EmitRunComplete(ctx, result.Success, result.TotalSteps, result.TotalTokens, result.Duration.Milliseconds(), result.FinishReason, err) + } +} + +func (e *AgenticExecutor) emitStepComplete(ctx context.Context, stepNum int, totalTokens int, resp *ai.ModelResponse) { + if e.streamCtx == nil { + return + } + inputTokens := 0 + outputTokens := 0 + if resp.Usage != nil { + inputTokens = resp.Usage.InputTokens + outputTokens = resp.Usage.OutputTokens + } + finishReason := "tool_use" + if len(resp.ToolRequests()) == 0 { + finishReason = "stop" + } + e.streamCtx.EmitStepComplete(ctx, stepNum, totalTokens, inputTokens, outputTokens, finishReason) +} + +func (e *AgenticExecutor) setup(ctx context.Context, agentID string, task string) error { + _, span := tracer.Start(ctx, "harness_setup") + defer span.End() + + if err := e.initializeSandbox(ctx); err != nil { + span.RecordError(err) + return fmt.Errorf("sandbox init failed: %w", err) + } + + if e.workspace != nil { + _, wsSpan := tracer.Start(ctx, "workspace_init") + if err := e.workspace.Initialize(ctx); err != nil { + wsSpan.RecordError(err) + wsSpan.End() + return fmt.Errorf("workspace init failed: %w", err) + } + wsSpan.SetAttributes( + attribute.String("workspace_path", e.workspace.Path()), + attribute.String("workspace_mode", e.config.Workspace.Mode), + ) + wsSpan.End() + } + + if e.gitManager != nil && e.config.Git.AutoBranch { + _, gitSpan := tracer.Start(ctx, "git_setup") + branchName, err := e.gitManager.CreateBranch(ctx, task, agentID) + if err != nil { + gitSpan.RecordError(err) + gitSpan.End() + return fmt.Errorf("git branch creation failed: %w", err) + } + gitSpan.SetAttributes(attribute.String("branch", branchName)) + gitSpan.End() + } + + return nil +} + +func (e *AgenticExecutor) initializeSandbox(ctx context.Context) error { + if e.sandbox != nil { + return nil + } + + if e.sandboxConfig == nil { + return nil + } + + _, span := tracer.Start(ctx, "sandbox_init") + defer span.End() + + factory := sandbox.NewFactory(sandbox.DefaultConfig()) + sb, err := factory.Create(*e.sandboxConfig) + if err != nil { + span.RecordError(err) + return fmt.Errorf("failed to create sandbox: %w", err) + } + + if err := sb.Create(ctx); err != nil { + span.RecordError(err) + return fmt.Errorf("failed to initialize sandbox: %w", err) + } + + e.sandbox = sb + span.SetAttributes( + attribute.String("sandbox_mode", string(e.sandboxConfig.Mode)), + attribute.String("sandbox_id", sb.ID()), + ) + + return nil +} + +func (e *AgenticExecutor) cleanup(ctx context.Context, agentID string, result *ExecutionResult) error { + if e.gitManager != nil && e.config.Git.AutoCommit && result.Success { + _, span := tracer.Start(ctx, "git_commit") + commitSHA, err := e.gitManager.Commit(ctx, result.Response) + if err != nil { + span.RecordError(err) + span.End() + return err + } + span.SetAttributes(attribute.String("commit_sha", commitSHA)) + span.End() + } + + if err := e.destroySandbox(ctx); err != nil { + return fmt.Errorf("sandbox cleanup failed: %w", err) + } + + return nil +} + +func (e *AgenticExecutor) destroySandbox(ctx context.Context) error { + if e.sandbox == nil { + return nil + } + + _, span := tracer.Start(ctx, "sandbox_destroy") + defer span.End() + + sandboxID := e.sandbox.ID() + if err := e.sandbox.Destroy(ctx); err != nil { + span.RecordError(err) + return err + } + + span.SetAttributes(attribute.String("sandbox_id", sandboxID)) + e.sandbox = nil + return nil +} + +func (e *AgenticExecutor) Sandbox() sandbox.Sandbox { + return e.sandbox +} + +func (e *AgenticExecutor) runLoop( + ctx context.Context, + task string, + tools []ai.Tool, +) (*ExecutionResult, error) { + // Initialize history from previous session if provided + var history []*ai.Message + if len(e.initialHistory) > 0 { + history = make([]*ai.Message, len(e.initialHistory)) + copy(history, e.initialHistory) + } + totalTokens := 0 + + toolRefs := make([]ai.ToolRef, len(tools)) + for i, t := range tools { + toolRefs[i] = t + } + + toolMap := make(map[string]ai.Tool) + for _, t := range tools { + toolMap[t.Name()] = t + } + + for step := 1; step <= e.agentConfig.MaxSteps; step++ { + select { + case <-ctx.Done(): + return &ExecutionResult{ + Success: false, + Error: "context cancelled", + TotalSteps: step - 1, + TotalTokens: totalTokens, + FinishReason: "cancelled", + History: history, + }, ctx.Err() + default: + } + + stepCtx, stepSpan := tracer.Start(ctx, "step", + trace.WithAttributes(attribute.Int("step", step)), + ) + + if e.compactor != nil && len(history) > 0 { + compactedHistory, compacted, err := e.compactor.CompactIfNeeded(stepCtx, history) + if err != nil { + stepSpan.RecordError(err) + } else if compacted { + history = compactedHistory + stepSpan.AddEvent("compaction_applied") + } + } + + genCtx, genSpan := tracer.Start(stepCtx, "llm_generate") + + // Build enhanced system prompt with skills injection + enhancedPrompt := e.buildEnhancedSystemPrompt() + + generateOpts := []ai.GenerateOption{ + ai.WithModelName(e.modelName), + ai.WithMessages(history...), + ai.WithTools(toolRefs...), + } + + if enhancedPrompt != "" { + generateOpts = append(generateOpts, ai.WithSystem(enhancedPrompt)) + } + + if step == 1 { + generateOpts = append(generateOpts, ai.WithPrompt(task)) + } + + resp, err := genkit.Generate(genCtx, e.genkitApp, generateOpts...) + + if err != nil { + genSpan.RecordError(err) + genSpan.End() + stepSpan.End() + return &ExecutionResult{ + Success: false, + Error: fmt.Sprintf("generation failed: %v", err), + TotalSteps: step, + TotalTokens: totalTokens, + FinishReason: "error", + History: history, + }, err + } + + if resp.Usage != nil { + totalTokens += resp.Usage.InputTokens + resp.Usage.OutputTokens + genSpan.SetAttributes( + attribute.Int("input_tokens", resp.Usage.InputTokens), + attribute.Int("output_tokens", resp.Usage.OutputTokens), + ) + } + genSpan.End() + + toolReqs := resp.ToolRequests() + if len(toolReqs) == 0 { + // Add final response to history for persistence + history = append(history, resp.Message) + stepSpan.SetAttributes(attribute.String("finish_reason", "agent_done")) + stepSpan.End() + return &ExecutionResult{ + Success: true, + Response: resp.Text(), + TotalSteps: step, + TotalTokens: totalTokens, + FinishReason: "agent_done", + History: history, + }, nil + } + + var toolParts []*ai.Part + for _, req := range toolReqs { + toolResult, toolErr := e.executeTool(stepCtx, req, toolMap) + if toolErr != nil { + toolParts = append(toolParts, ai.NewToolResponsePart(&ai.ToolResponse{ + Name: req.Name, + Ref: req.Ref, + Output: map[string]string{"error": toolErr.Error()}, + })) + continue + } + toolParts = append(toolParts, ai.NewToolResponsePart(&ai.ToolResponse{ + Name: req.Name, + Ref: req.Ref, + Output: toolResult, + })) + } + + history = append(history, resp.Message) + history = append(history, ai.NewMessage(ai.RoleTool, nil, toolParts...)) + + e.emitStepComplete(stepCtx, step, totalTokens, resp) + + stepSpan.SetAttributes(attribute.Int("tool_count", len(toolReqs))) + stepSpan.End() + } + + return &ExecutionResult{ + Success: false, + Error: "max steps exceeded", + TotalSteps: e.agentConfig.MaxSteps, + TotalTokens: totalTokens, + FinishReason: "max_steps", + History: history, + }, nil +} + +func (e *AgenticExecutor) executeTool( + ctx context.Context, + req *ai.ToolRequest, + toolMap map[string]ai.Tool, +) (interface{}, error) { + _, span := tracer.Start(ctx, "tool_execution", + trace.WithAttributes( + attribute.String("tool", req.Name), + ), + ) + defer span.End() + + toolID := req.Ref + if toolID == "" { + toolID = fmt.Sprintf("%s-%d", req.Name, time.Now().UnixNano()) + } + + e.emitToolStart(ctx, req.Name, toolID, req.Input) + + preCtx, preSpan := tracer.Start(ctx, "pre_tool_hook") + hookResult, hookMsg := e.hooks.RunPreHooks(preCtx, req) + preSpan.SetAttributes( + attribute.String("result", string(hookResult)), + ) + preSpan.End() + + switch hookResult { + case HookBlock: + span.SetAttributes(attribute.Bool("blocked", true)) + err := fmt.Errorf("tool blocked: %s", hookMsg) + e.emitToolResult(ctx, req.Name, toolID, nil, 0, err) + return nil, err + case HookInterrupt: + span.SetAttributes(attribute.Bool("interrupted", true)) + err := fmt.Errorf("tool requires approval: %s", hookMsg) + e.emitToolResult(ctx, req.Name, toolID, nil, 0, err) + return nil, err + } + + tool, exists := toolMap[req.Name] + if !exists { + err := fmt.Errorf("tool not found: %s", req.Name) + e.emitToolResult(ctx, req.Name, toolID, nil, 0, err) + return nil, err + } + + runCtx, runSpan := tracer.Start(ctx, "tool_run") + startTime := time.Now() + + output, err := tool.RunRaw(runCtx, req.Input) + + durationMs := time.Since(startTime).Milliseconds() + runSpan.SetAttributes( + attribute.Int64("duration_ms", durationMs), + attribute.Bool("success", err == nil), + ) + runSpan.End() + + e.emitToolResult(ctx, req.Name, toolID, output, durationMs, err) + + if err != nil { + span.RecordError(err) + return nil, err + } + + postCtx, postSpan := tracer.Start(ctx, "post_tool_hook") + e.hooks.RunPostHooks(postCtx, req, output) + postSpan.End() + + e.doomDetector.Record(req.Name, req.Input) + + return output, nil +} + +func (e *AgenticExecutor) emitToolStart(ctx context.Context, toolName, toolID string, input any) { + if e.streamCtx != nil { + e.streamCtx.EmitToolStart(ctx, toolName, toolID, input) + } +} + +func (e *AgenticExecutor) emitToolResult(ctx context.Context, toolName, toolID string, output any, durationMs int64, err error) { + if e.streamCtx != nil { + e.streamCtx.EmitToolResult(ctx, toolName, toolID, output, durationMs, err) + } +} + +// buildEnhancedSystemPrompt combines the base system prompt with skills and memory sections +func (e *AgenticExecutor) buildEnhancedSystemPrompt() string { + var parts []string + + // Base agent prompt + if e.systemPrompt != "" { + parts = append(parts, e.systemPrompt) + } + + // Skills injection (progressive disclosure - only names and descriptions) + if e.skillsMiddleware != nil { + skillsList, err := e.skillsMiddleware.LoadSkills() + if err == nil && len(skillsList) > 0 { + skillsSection := e.skillsMiddleware.FormatSystemPromptSection(skillsList) + if skillsSection != "" { + parts = append(parts, skillsSection) + } + } + } + + // Memory injection (always loaded, unlike skills) + if e.memoryMiddleware != nil { + memorySection, err := e.memoryMiddleware.FormatSystemPromptSection() + if err == nil && memorySection != "" { + parts = append(parts, memorySection) + } + } + + return strings.Join(parts, "\n") +} + +type WorkspaceManager interface { + Initialize(ctx context.Context) error + Path() string + ReadFile(ctx context.Context, path string) ([]byte, error) + WriteFile(ctx context.Context, path string, data []byte) error + Close(ctx context.Context) error +} + +type GitManager interface { + CreateBranch(ctx context.Context, task string, agentID string) (string, error) + Commit(ctx context.Context, message string) (string, error) + Push(ctx context.Context) error + GetCurrentBranch(ctx context.Context) (string, error) +} + +type PromptBuilder struct { + config *HarnessConfig +} diff --git a/pkg/harness/executor_e2e_test.go b/pkg/harness/executor_e2e_test.go new file mode 100644 index 00000000..e323df6e --- /dev/null +++ b/pkg/harness/executor_e2e_test.go @@ -0,0 +1,777 @@ +package harness_test + +import ( + "context" + "fmt" + "os" + "os/exec" + "strings" + "testing" + "time" + + "station/internal/config" + "station/internal/genkit/anthropic_oauth" + "station/internal/logging" + "station/pkg/harness" + "station/pkg/harness/sandbox" + "station/pkg/harness/tools" + "station/pkg/harness/workspace" + + "github.com/firebase/genkit/go/genkit" + "github.com/firebase/genkit/go/plugins/compat_oai/openai" + "github.com/firebase/genkit/go/plugins/googlegenai" +) + +func TestAgenticExecutor_E2E_RealLLM(t *testing.T) { + if os.Getenv("HARNESS_E2E_TEST") != "1" { + t.Skip("Skipping e2e test. Set HARNESS_E2E_TEST=1 to run with real LLM") + } + + ctx := context.Background() + + configDir := config.GetStationConfigDir() + configPath := configDir + "/config.yaml" + t.Logf("Looking for config in: %s", configDir) + + if _, err := os.Stat(configPath); err != nil { + t.Logf("Config file not found at %s: %v", configPath, err) + } else { + t.Logf("Config file exists at %s", configPath) + } + + if err := config.InitViper(""); err != nil { + t.Fatalf("Failed to init viper: %v", err) + } + + cfg, err := config.Load() + if err != nil { + t.Fatalf("Failed to load config: %v", err) + } + + t.Logf("Loaded config - AIProvider: '%s', AIModel: '%s', AIAuthType: '%s', OAuthToken length: %d", + cfg.AIProvider, cfg.AIModel, cfg.AIAuthType, len(cfg.AIOAuthToken)) + + genkitApp, err := initGenkit(ctx, cfg) + if err != nil { + t.Fatalf("Failed to initialize GenKit: %v", err) + } + + tmpDir, err := os.MkdirTemp("", "harness-e2e-*") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + defer os.RemoveAll(tmpDir) + + ws := workspace.NewHostWorkspace(tmpDir) + if err := ws.Initialize(ctx); err != nil { + t.Fatalf("Failed to initialize workspace: %v", err) + } + + toolRegistry := tools.NewToolRegistry(genkitApp, tmpDir) + if err := toolRegistry.RegisterBuiltinTools(); err != nil { + t.Fatalf("Failed to register tools: %v", err) + } + + harnessConfig := harness.DefaultHarnessConfig() + agentConfig := &harness.AgentHarnessConfig{ + MaxSteps: 10, + DoomLoopThreshold: 3, + Timeout: 2 * time.Minute, + } + + modelName := formatModelName(cfg.AIProvider, cfg.AIModel) + + executor := harness.NewAgenticExecutor( + genkitApp, + harnessConfig, + agentConfig, + harness.WithWorkspace(ws), + harness.WithModelName(modelName), + ) + + task := "Create a file called 'hello.txt' with the content 'Hello from the agentic harness!' then read it back and confirm the content." + + t.Logf("Executing task: %s", task) + startTime := time.Now() + + result, err := executor.Execute(ctx, "e2e-test", task, toolRegistry.All()) + + t.Logf("Execution completed in %v", time.Since(startTime)) + t.Logf("Result: success=%v, steps=%d, tokens=%d, finish_reason=%s", + result.Success, result.TotalSteps, result.TotalTokens, result.FinishReason) + + if err != nil { + t.Logf("Execution error: %v", err) + } + + if result.Response != "" { + t.Logf("Response:\n%s", result.Response) + } + + if !result.Success { + t.Errorf("Expected successful execution, got error: %s", result.Error) + } + + helloPath := tmpDir + "/hello.txt" + if _, err := os.Stat(helloPath); os.IsNotExist(err) { + t.Error("Expected hello.txt to be created") + } else { + content, _ := os.ReadFile(helloPath) + t.Logf("Created file content: %s", string(content)) + } +} + +func TestAgenticExecutor_E2E_MultiStep(t *testing.T) { + if os.Getenv("HARNESS_E2E_TEST") != "1" { + t.Skip("Skipping e2e test. Set HARNESS_E2E_TEST=1 to run with real LLM") + } + + ctx := context.Background() + + if err := config.InitViper(""); err != nil { + t.Fatalf("Failed to init viper: %v", err) + } + + cfg, err := config.Load() + if err != nil { + t.Fatalf("Failed to load config: %v", err) + } + + genkitApp, err := initGenkit(ctx, cfg) + if err != nil { + t.Fatalf("Failed to initialize GenKit: %v", err) + } + + tmpDir, err := os.MkdirTemp("", "harness-e2e-multi-*") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + defer os.RemoveAll(tmpDir) + + ws := workspace.NewHostWorkspace(tmpDir) + if err := ws.Initialize(ctx); err != nil { + t.Fatalf("Failed to initialize workspace: %v", err) + } + + toolRegistry := tools.NewToolRegistry(genkitApp, tmpDir) + if err := toolRegistry.RegisterBuiltinTools(); err != nil { + t.Fatalf("Failed to register tools: %v", err) + } + + harnessConfig := harness.DefaultHarnessConfig() + agentConfig := &harness.AgentHarnessConfig{ + MaxSteps: 15, + DoomLoopThreshold: 3, + Timeout: 3 * time.Minute, + } + + modelName := formatModelName(cfg.AIProvider, cfg.AIModel) + + executor := harness.NewAgenticExecutor( + genkitApp, + harnessConfig, + agentConfig, + harness.WithWorkspace(ws), + harness.WithModelName(modelName), + ) + + task := `Do the following steps: +1. Create a directory called 'project' +2. Inside 'project', create a file called 'config.json' with content: {"name": "test", "version": "1.0"} +3. List all files in the workspace to confirm +4. Read the config.json file and tell me what version it contains` + + t.Logf("Executing multi-step task") + startTime := time.Now() + + result, err := executor.Execute(ctx, "e2e-multi-test", task, toolRegistry.All()) + + t.Logf("Execution completed in %v", time.Since(startTime)) + t.Logf("Result: success=%v, steps=%d, tokens=%d, finish_reason=%s", + result.Success, result.TotalSteps, result.TotalTokens, result.FinishReason) + + if err != nil { + t.Logf("Execution error: %v", err) + } + + if result.Response != "" { + t.Logf("Response:\n%s", result.Response) + } + + if !result.Success { + t.Errorf("Expected successful execution, got error: %s", result.Error) + } + + configPath := tmpDir + "/project/config.json" + if _, err := os.Stat(configPath); os.IsNotExist(err) { + t.Error("Expected project/config.json to be created") + } +} + +func initGenkit(ctx context.Context, cfg *config.Config) (*genkit.Genkit, error) { + os.Setenv("OTEL_SDK_DISABLED", "true") + + promptDir := "/tmp/harness-e2e-prompts" + os.MkdirAll(promptDir, 0755) + + provider := strings.ToLower(cfg.AIProvider) + + if strings.HasPrefix(strings.ToLower(cfg.AIModel), "claude") { + provider = "anthropic" + } else if strings.HasPrefix(strings.ToLower(cfg.AIModel), "gemini") { + provider = "gemini" + } else if strings.HasPrefix(strings.ToLower(cfg.AIModel), "gpt") { + provider = "openai" + } + + logging.Info("Initializing GenKit for e2e test with provider: %s, model: %s", provider, cfg.AIModel) + + switch provider { + case "openai": + return genkit.Init(ctx, + genkit.WithPlugins(&openai.OpenAI{APIKey: cfg.AIAPIKey}), + genkit.WithPromptDir(promptDir)), nil + + case "gemini", "googlegenai": + return genkit.Init(ctx, + genkit.WithPlugins(&googlegenai.GoogleAI{}), + genkit.WithPromptDir(promptDir)), nil + + case "anthropic": + var plugin *anthropic_oauth.AnthropicOAuth + if cfg.AIAuthType == "oauth" && cfg.AIOAuthToken != "" { + plugin = &anthropic_oauth.AnthropicOAuth{OAuthToken: cfg.AIOAuthToken} + } else if cfg.AIAPIKey != "" { + plugin = &anthropic_oauth.AnthropicOAuth{APIKey: cfg.AIAPIKey} + } else { + return nil, fmt.Errorf("Anthropic requires OAuth token or API key") + } + return genkit.Init(ctx, + genkit.WithPlugins(plugin), + genkit.WithPromptDir(promptDir)), nil + + default: + return nil, fmt.Errorf("unsupported provider: %s", provider) + } +} + +func formatModelName(provider, model string) string { + provider = strings.ToLower(provider) + + if strings.HasPrefix(strings.ToLower(model), "claude") { + provider = "anthropic" + } else if strings.HasPrefix(strings.ToLower(model), "gemini") { + provider = "googlegenai" + } else if strings.HasPrefix(strings.ToLower(model), "gpt") { + provider = "openai" + } + + return provider + "/" + model +} + +func TestAgenticExecutor_E2E_Compaction(t *testing.T) { + if os.Getenv("HARNESS_E2E_TEST") != "1" { + t.Skip("Skipping e2e test. Set HARNESS_E2E_TEST=1 to run with real LLM") + } + + ctx := context.Background() + + if err := config.InitViper(""); err != nil { + t.Fatalf("Failed to init viper: %v", err) + } + + cfg, err := config.Load() + if err != nil { + t.Fatalf("Failed to load config: %v", err) + } + + genkitApp, err := initGenkit(ctx, cfg) + if err != nil { + t.Fatalf("Failed to initialize GenKit: %v", err) + } + + tmpDir, err := os.MkdirTemp("", "harness-e2e-compaction-*") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + defer os.RemoveAll(tmpDir) + + ws := workspace.NewHostWorkspace(tmpDir) + if err := ws.Initialize(ctx); err != nil { + t.Fatalf("Failed to initialize workspace: %v", err) + } + + toolRegistry := tools.NewToolRegistry(genkitApp, tmpDir) + if err := toolRegistry.RegisterBuiltinTools(); err != nil { + t.Fatalf("Failed to register tools: %v", err) + } + + harnessConfig := harness.DefaultHarnessConfig() + harnessConfig.Compaction.Enabled = true + harnessConfig.Compaction.Threshold = 0.3 + harnessConfig.Compaction.ProtectTokens = 200 + + agentConfig := &harness.AgentHarnessConfig{ + MaxSteps: 10, + DoomLoopThreshold: 5, + Timeout: 3 * time.Minute, + } + + modelName := formatModelName(cfg.AIProvider, cfg.AIModel) + + compactor := harness.NewCompactor(genkitApp, modelName, harnessConfig.Compaction, 1000) + + executor := harness.NewAgenticExecutor( + genkitApp, + harnessConfig, + agentConfig, + harness.WithWorkspace(ws), + harness.WithModelName(modelName), + harness.WithCompactor(compactor), + ) + + task := `Create two files: 'data.txt' with content 'Important data here' and 'config.json' with content '{"version": "1.0"}'. Then read both files back and summarize what they contain.` + + t.Logf("Executing compaction test task with context window=1000, threshold=0.3") + startTime := time.Now() + + result, err := executor.Execute(ctx, "e2e-compaction-test", task, toolRegistry.All()) + + t.Logf("Execution completed in %v", time.Since(startTime)) + t.Logf("Result: success=%v, steps=%d, tokens=%d, finish_reason=%s", + result.Success, result.TotalSteps, result.TotalTokens, result.FinishReason) + + if err != nil { + t.Logf("Execution error: %v", err) + } + + if result.Response != "" { + t.Logf("Response:\n%s", result.Response) + } + + if !result.Success { + t.Errorf("Expected successful execution, got error: %s", result.Error) + } + + dataPath := tmpDir + "/data.txt" + if _, err := os.Stat(dataPath); os.IsNotExist(err) { + t.Error("Expected data.txt to be created") + } + + configPath := tmpDir + "/config.json" + if _, err := os.Stat(configPath); os.IsNotExist(err) { + t.Error("Expected config.json to be created") + } + + t.Logf("Compaction test completed - context window was limited to 1000 tokens with 0.3 threshold") +} + +func TestAgenticExecutor_E2E_WorkflowSimulation(t *testing.T) { + if os.Getenv("HARNESS_E2E_TEST") != "1" { + t.Skip("Skipping e2e test. Set HARNESS_E2E_TEST=1 to run with real LLM") + } + + ctx := context.Background() + + if err := config.InitViper(""); err != nil { + t.Fatalf("Failed to init viper: %v", err) + } + + cfg, err := config.Load() + if err != nil { + t.Fatalf("Failed to load config: %v", err) + } + + genkitApp, err := initGenkit(ctx, cfg) + if err != nil { + t.Fatalf("Failed to initialize GenKit: %v", err) + } + + tmpDir, err := os.MkdirTemp("", "harness-e2e-workflow-*") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + defer os.RemoveAll(tmpDir) + + ws := workspace.NewHostWorkspace(tmpDir) + if err := ws.Initialize(ctx); err != nil { + t.Fatalf("Failed to initialize workspace: %v", err) + } + + toolRegistry := tools.NewToolRegistry(genkitApp, tmpDir) + if err := toolRegistry.RegisterBuiltinTools(); err != nil { + t.Fatalf("Failed to register tools: %v", err) + } + + harnessConfig := harness.DefaultHarnessConfig() + modelName := formatModelName(cfg.AIProvider, cfg.AIModel) + + t.Log("=== Simulating Multi-Agent Workflow with Harness ===") + t.Log("Step 1: Setup Agent - Create project structure") + + setupConfig := &harness.AgentHarnessConfig{ + MaxSteps: 10, + DoomLoopThreshold: 3, + Timeout: 2 * time.Minute, + } + + setupExecutor := harness.NewAgenticExecutor( + genkitApp, + harnessConfig, + setupConfig, + harness.WithWorkspace(ws), + harness.WithModelName(modelName), + ) + + setupTask := "Create a directory called 'test-project' and inside it create a file called 'config.json' with content '{\"name\": \"test\", \"version\": \"1.0\"}'" + + setupResult, err := setupExecutor.Execute(ctx, "setup-agent", setupTask, toolRegistry.All()) + if err != nil { + t.Fatalf("Setup agent failed: %v", err) + } + + t.Logf("Setup agent: success=%v, steps=%d, tokens=%d", setupResult.Success, setupResult.TotalSteps, setupResult.TotalTokens) + + if !setupResult.Success { + t.Fatalf("Setup agent should succeed: %s", setupResult.Error) + } + + configPath := tmpDir + "/test-project/config.json" + if _, err := os.Stat(configPath); os.IsNotExist(err) { + t.Fatal("Setup agent should have created test-project/config.json") + } + + t.Log("Step 2: Verify Agent - Check created files") + + verifyConfig := &harness.AgentHarnessConfig{ + MaxSteps: 10, + DoomLoopThreshold: 3, + Timeout: 2 * time.Minute, + } + + verifyExecutor := harness.NewAgenticExecutor( + genkitApp, + harnessConfig, + verifyConfig, + harness.WithWorkspace(ws), + harness.WithModelName(modelName), + ) + + verifyTask := "List all files in the 'test-project' directory and read config.json. Report the version number found." + + verifyResult, err := verifyExecutor.Execute(ctx, "verify-agent", verifyTask, toolRegistry.All()) + if err != nil { + t.Fatalf("Verify agent failed: %v", err) + } + + t.Logf("Verify agent: success=%v, steps=%d, tokens=%d", verifyResult.Success, verifyResult.TotalSteps, verifyResult.TotalTokens) + t.Logf("Verify response: %s", verifyResult.Response) + + if !verifyResult.Success { + t.Fatalf("Verify agent should succeed: %s", verifyResult.Error) + } + + if !strings.Contains(verifyResult.Response, "1.0") { + t.Error("Verify agent should report version 1.0") + } + + t.Log("=== Workflow Simulation Complete ===") + t.Logf("Total steps: Setup=%d, Verify=%d", setupResult.TotalSteps, verifyResult.TotalSteps) + t.Logf("Total tokens: Setup=%d, Verify=%d", setupResult.TotalTokens, verifyResult.TotalTokens) +} + +func TestAgenticExecutor_E2E_DockerSandbox(t *testing.T) { + if os.Getenv("HARNESS_E2E_TEST") != "1" { + t.Skip("Skipping e2e test. Set HARNESS_E2E_TEST=1 to run with real LLM") + } + + if _, err := exec.LookPath("docker"); err != nil { + t.Skip("Docker not available, skipping Docker sandbox test") + } + + ctx := context.Background() + + if err := config.InitViper(""); err != nil { + t.Fatalf("Failed to init viper: %v", err) + } + + cfg, err := config.Load() + if err != nil { + t.Fatalf("Failed to load config: %v", err) + } + + genkitApp, err := initGenkit(ctx, cfg) + if err != nil { + t.Fatalf("Failed to initialize GenKit: %v", err) + } + + tmpDir, err := os.MkdirTemp("", "harness-e2e-docker-*") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + defer os.RemoveAll(tmpDir) + + t.Logf("Using workspace directory: %s", tmpDir) + + sbConfig := sandbox.Config{ + Mode: sandbox.ModeDocker, + Image: "python:3.11-slim", + WorkspacePath: tmpDir, + Timeout: 5 * time.Minute, + Resources: sandbox.ResourceConfig{ + CPU: 1, + Memory: "512m", + PIDs: 100, + }, + Network: sandbox.NetworkConfig{ + Enabled: false, + }, + } + + factory := sandbox.NewFactory(sandbox.DefaultConfig()) + sb, err := factory.Create(sbConfig) + if err != nil { + t.Fatalf("Failed to create Docker sandbox: %v", err) + } + defer func() { + if err := sb.Destroy(ctx); err != nil { + t.Logf("Warning: failed to destroy sandbox: %v", err) + } + }() + + if err := sb.Create(ctx); err != nil { + t.Fatalf("Failed to initialize Docker sandbox: %v", err) + } + + t.Logf("Docker sandbox created: id=%s, image=%s", sb.ID(), sbConfig.Image) + + ws := workspace.NewHostWorkspace(tmpDir) + if err := ws.Initialize(ctx); err != nil { + t.Fatalf("Failed to initialize workspace: %v", err) + } + + toolRegistry := tools.NewToolRegistryWithSandbox(genkitApp, tmpDir, sb) + if err := toolRegistry.RegisterBuiltinTools(); err != nil { + t.Fatalf("Failed to register tools: %v", err) + } + + harnessConfig := harness.DefaultHarnessConfig() + agentConfig := &harness.AgentHarnessConfig{ + MaxSteps: 15, + DoomLoopThreshold: 3, + Timeout: 3 * time.Minute, + } + + modelName := formatModelName(cfg.AIProvider, cfg.AIModel) + + executor := harness.NewAgenticExecutor( + genkitApp, + harnessConfig, + agentConfig, + harness.WithWorkspace(ws), + harness.WithModelName(modelName), + harness.WithSandbox(sb), + ) + + task := `Create a Python script called 'hello.py' with this content: +print("Hello from Docker sandbox!") +import sys +print(f"Python version: {sys.version}") + +Then run it using: python hello.py + +Report the output.` + + t.Logf("Executing task in Docker sandbox...") + startTime := time.Now() + + result, err := executor.Execute(ctx, "docker-e2e-test", task, toolRegistry.All()) + + t.Logf("Execution completed in %v", time.Since(startTime)) + t.Logf("Result: success=%v, steps=%d, tokens=%d, finish_reason=%s", + result.Success, result.TotalSteps, result.TotalTokens, result.FinishReason) + + if err != nil { + t.Logf("Execution error: %v", err) + } + + if result.Response != "" { + t.Logf("Response:\n%s", result.Response) + } + + if !result.Success { + t.Errorf("Expected successful execution, got error: %s", result.Error) + } + + scriptPath := tmpDir + "/hello.py" + if _, err := os.Stat(scriptPath); os.IsNotExist(err) { + t.Error("Expected hello.py to be created on host (via mounted volume)") + } else { + content, _ := os.ReadFile(scriptPath) + t.Logf("Script content on host:\n%s", string(content)) + } + + if !strings.Contains(result.Response, "Hello from Docker sandbox") && + !strings.Contains(result.Response, "Python version") { + t.Logf("Warning: Response may not contain expected output") + } +} + +func TestAgenticExecutor_E2E_DockerWorkflowHandoff(t *testing.T) { + if os.Getenv("HARNESS_E2E_TEST") != "1" { + t.Skip("Skipping e2e test. Set HARNESS_E2E_TEST=1 to run with real LLM") + } + + if _, err := exec.LookPath("docker"); err != nil { + t.Skip("Docker not available, skipping Docker sandbox test") + } + + ctx := context.Background() + + if err := config.InitViper(""); err != nil { + t.Fatalf("Failed to init viper: %v", err) + } + + cfg, err := config.Load() + if err != nil { + t.Fatalf("Failed to load config: %v", err) + } + + genkitApp, err := initGenkit(ctx, cfg) + if err != nil { + t.Fatalf("Failed to initialize GenKit: %v", err) + } + + tmpDir, err := os.MkdirTemp("", "harness-e2e-docker-handoff-*") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + defer os.RemoveAll(tmpDir) + + t.Log("=== Docker Workflow Handoff Test ===") + t.Logf("Shared workspace: %s", tmpDir) + + harnessConfig := harness.DefaultHarnessConfig() + modelName := formatModelName(cfg.AIProvider, cfg.AIModel) + + t.Log("Step 1: Agent A creates files in Docker container") + + sbConfigA := sandbox.Config{ + Mode: sandbox.ModeDocker, + Image: "python:3.11-slim", + WorkspacePath: tmpDir, + Timeout: 3 * time.Minute, + Resources: sandbox.ResourceConfig{ + CPU: 1, + Memory: "512m", + }, + } + + factoryA := sandbox.NewFactory(sandbox.DefaultConfig()) + sbA, err := factoryA.Create(sbConfigA) + if err != nil { + t.Fatalf("Failed to create sandbox A: %v", err) + } + + if err := sbA.Create(ctx); err != nil { + t.Fatalf("Failed to initialize sandbox A: %v", err) + } + + t.Logf("Sandbox A created: %s", sbA.ID()) + + wsA := workspace.NewHostWorkspace(tmpDir) + wsA.Initialize(ctx) + + toolRegistryA := tools.NewToolRegistryWithSandbox(genkitApp, tmpDir, sbA) + toolRegistryA.RegisterBuiltinTools() + + executorA := harness.NewAgenticExecutor( + genkitApp, + harnessConfig, + &harness.AgentHarnessConfig{MaxSteps: 10, DoomLoopThreshold: 3, Timeout: 2 * time.Minute}, + harness.WithWorkspace(wsA), + harness.WithModelName(modelName), + harness.WithSandbox(sbA), + ) + + taskA := "Create a file called 'data.json' with content: {\"created_by\": \"agent_a\", \"value\": 42}" + + resultA, err := executorA.Execute(ctx, "agent-a", taskA, toolRegistryA.All()) + if err != nil || !resultA.Success { + t.Fatalf("Agent A failed: %v / %s", err, resultA.Error) + } + t.Logf("Agent A completed: steps=%d", resultA.TotalSteps) + + if err := sbA.Destroy(ctx); err != nil { + t.Logf("Warning: failed to destroy sandbox A: %v", err) + } + t.Log("Sandbox A destroyed") + + dataPath := tmpDir + "/data.json" + if _, err := os.Stat(dataPath); os.IsNotExist(err) { + t.Fatal("data.json should exist on host after Agent A") + } + contentAfterA, _ := os.ReadFile(dataPath) + t.Logf("File content after Agent A: %s", string(contentAfterA)) + + t.Log("Step 2: Agent B reads and modifies files in NEW Docker container") + + sbConfigB := sandbox.Config{ + Mode: sandbox.ModeDocker, + Image: "python:3.11-slim", + WorkspacePath: tmpDir, + Timeout: 3 * time.Minute, + Resources: sandbox.ResourceConfig{ + CPU: 1, + Memory: "512m", + }, + } + + factoryB := sandbox.NewFactory(sandbox.DefaultConfig()) + sbB, err := factoryB.Create(sbConfigB) + if err != nil { + t.Fatalf("Failed to create sandbox B: %v", err) + } + + if err := sbB.Create(ctx); err != nil { + t.Fatalf("Failed to initialize sandbox B: %v", err) + } + + t.Logf("Sandbox B created: %s (different container!)", sbB.ID()) + + wsB := workspace.NewHostWorkspace(tmpDir) + wsB.Initialize(ctx) + + toolRegistryB := tools.NewToolRegistryWithSandbox(genkitApp, tmpDir, sbB) + toolRegistryB.RegisterBuiltinTools() + + executorB := harness.NewAgenticExecutor( + genkitApp, + harnessConfig, + &harness.AgentHarnessConfig{MaxSteps: 10, DoomLoopThreshold: 3, Timeout: 2 * time.Minute}, + harness.WithWorkspace(wsB), + harness.WithModelName(modelName), + harness.WithSandbox(sbB), + ) + + taskB := "Read the file 'data.json', report who created it and what value is stored." + + resultB, err := executorB.Execute(ctx, "agent-b", taskB, toolRegistryB.All()) + if err != nil || !resultB.Success { + t.Fatalf("Agent B failed: %v / %s", err, resultB.Error) + } + t.Logf("Agent B completed: steps=%d", resultB.TotalSteps) + t.Logf("Agent B response: %s", resultB.Response) + + if err := sbB.Destroy(ctx); err != nil { + t.Logf("Warning: failed to destroy sandbox B: %v", err) + } + + if !strings.Contains(resultB.Response, "agent_a") || !strings.Contains(resultB.Response, "42") { + t.Error("Agent B should report data created by Agent A") + } + + t.Log("=== Docker Workflow Handoff Test Complete ===") + t.Log("Files created by Agent A were successfully read by Agent B in a DIFFERENT container") + t.Log("This proves workspace volume mounting enables session persistence across container destroys") +} diff --git a/pkg/harness/executor_test.go b/pkg/harness/executor_test.go new file mode 100644 index 00000000..2a9a10f5 --- /dev/null +++ b/pkg/harness/executor_test.go @@ -0,0 +1,157 @@ +package harness + +import ( + "context" + "testing" + + "station/pkg/harness/sandbox" +) + +func TestNewAgenticExecutor_WithSandboxOptions(t *testing.T) { + cfg := &sandbox.Config{ + Mode: sandbox.ModeHost, + Timeout: sandbox.DefaultConfig().Timeout, + } + + executor := NewAgenticExecutor( + nil, + nil, + nil, + WithSandboxConfig(cfg), + ) + + if executor.sandboxConfig == nil { + t.Error("sandboxConfig should be set") + } + if executor.sandboxConfig.Mode != sandbox.ModeHost { + t.Errorf("sandboxConfig.Mode = %s, want %s", executor.sandboxConfig.Mode, sandbox.ModeHost) + } +} + +func TestAgenticExecutor_InitializeSandbox_NoConfig(t *testing.T) { + executor := NewAgenticExecutor(nil, nil, nil) + + err := executor.initializeSandbox(context.Background()) + if err != nil { + t.Errorf("initializeSandbox should succeed with no config: %v", err) + } + if executor.sandbox != nil { + t.Error("sandbox should be nil when no config provided") + } +} + +func TestAgenticExecutor_InitializeSandbox_WithExistingSandbox(t *testing.T) { + hostSandbox, err := sandbox.NewHostSandbox(sandbox.DefaultConfig()) + if err != nil { + t.Fatalf("failed to create host sandbox: %v", err) + } + + executor := NewAgenticExecutor( + nil, + nil, + nil, + WithSandbox(hostSandbox), + ) + + err = executor.initializeSandbox(context.Background()) + if err != nil { + t.Errorf("initializeSandbox should succeed with existing sandbox: %v", err) + } + if executor.sandbox != hostSandbox { + t.Error("sandbox should remain unchanged when already provided") + } +} + +func TestAgenticExecutor_InitializeSandbox_CreatesHostSandbox(t *testing.T) { + cfg := &sandbox.Config{ + Mode: sandbox.ModeHost, + Timeout: sandbox.DefaultConfig().Timeout, + } + + executor := NewAgenticExecutor( + nil, + nil, + nil, + WithSandboxConfig(cfg), + ) + + err := executor.initializeSandbox(context.Background()) + if err != nil { + t.Errorf("initializeSandbox failed: %v", err) + } + if executor.sandbox == nil { + t.Error("sandbox should be created") + } + if executor.sandbox.ID() == "" { + t.Error("sandbox ID should not be empty") + } +} + +func TestAgenticExecutor_InitializeSandbox_CreatesDockerSandbox(t *testing.T) { + cfg := &sandbox.Config{ + Mode: sandbox.ModeDocker, + Image: "ubuntu:22.04", + Timeout: sandbox.DefaultConfig().Timeout, + } + + executor := NewAgenticExecutor( + nil, + nil, + nil, + WithSandboxConfig(cfg), + ) + + err := executor.initializeSandbox(context.Background()) + if err != nil { + t.Logf("Docker sandbox creation failed (expected if Docker not available): %v", err) + t.Skip("Skipping Docker test - Docker not available") + } + if executor.sandbox == nil { + t.Error("sandbox should be created") + } +} + +func TestAgenticExecutor_DestroySandbox(t *testing.T) { + hostSandbox, err := sandbox.NewHostSandbox(sandbox.DefaultConfig()) + if err != nil { + t.Fatalf("failed to create host sandbox: %v", err) + } + + executor := NewAgenticExecutor( + nil, + nil, + nil, + WithSandbox(hostSandbox), + ) + + err = executor.destroySandbox(context.Background()) + if err != nil { + t.Errorf("destroySandbox failed: %v", err) + } + if executor.sandbox != nil { + t.Error("sandbox should be nil after destroy") + } +} + +func TestAgenticExecutor_SandboxGetter(t *testing.T) { + hostSandbox, err := sandbox.NewHostSandbox(sandbox.DefaultConfig()) + if err != nil { + t.Fatalf("failed to create host sandbox: %v", err) + } + + executor := NewAgenticExecutor( + nil, + nil, + nil, + WithSandbox(hostSandbox), + ) + + if executor.Sandbox() != hostSandbox { + t.Error("Sandbox() should return the configured sandbox") + } + + executorNoSandbox := NewAgenticExecutor(nil, nil, nil) + if executorNoSandbox.Sandbox() != nil { + t.Error("Sandbox() should return nil when no sandbox configured") + } +} diff --git a/pkg/harness/git/cloner.go b/pkg/harness/git/cloner.go new file mode 100644 index 00000000..0cb69d39 --- /dev/null +++ b/pkg/harness/git/cloner.go @@ -0,0 +1,110 @@ +package git + +import ( + "context" + "fmt" + "os" + "os/exec" + "path/filepath" + "strings" +) + +type CloneConfig struct { + URL string + Branch string + Ref string + Depth int + Submodules bool + SSHKeyPath string +} + +type Cloner struct { + config CloneConfig +} + +func NewCloner(config CloneConfig) *Cloner { + return &Cloner{config: config} +} + +func (c *Cloner) Clone(ctx context.Context, targetPath string) error { + if c.config.URL == "" { + return fmt.Errorf("git URL is required") + } + + if err := os.MkdirAll(filepath.Dir(targetPath), 0755); err != nil { + return fmt.Errorf("failed to create parent directory: %w", err) + } + + args := []string{"clone"} + + if c.config.Depth > 0 { + args = append(args, "--depth", fmt.Sprintf("%d", c.config.Depth)) + } + + if c.config.Branch != "" && c.config.Ref == "" { + args = append(args, "--branch", c.config.Branch) + } + + args = append(args, c.config.URL, targetPath) + + cmd := exec.CommandContext(ctx, "git", args...) + cmd.Env = c.buildEnv() + + output, err := cmd.CombinedOutput() + if err != nil { + return fmt.Errorf("git clone failed: %s - %w", strings.TrimSpace(string(output)), err) + } + + if c.config.Ref != "" { + if err := c.checkoutRef(ctx, targetPath); err != nil { + return err + } + } + + if c.config.Submodules { + if err := c.initSubmodules(ctx, targetPath); err != nil { + return err + } + } + + return nil +} + +func (c *Cloner) checkoutRef(ctx context.Context, targetPath string) error { + cmd := exec.CommandContext(ctx, "git", "checkout", c.config.Ref) + cmd.Dir = targetPath + cmd.Env = c.buildEnv() + + output, err := cmd.CombinedOutput() + if err != nil { + return fmt.Errorf("git checkout %s failed: %s - %w", c.config.Ref, strings.TrimSpace(string(output)), err) + } + return nil +} + +func (c *Cloner) initSubmodules(ctx context.Context, targetPath string) error { + cmd := exec.CommandContext(ctx, "git", "submodule", "update", "--init", "--recursive") + cmd.Dir = targetPath + cmd.Env = c.buildEnv() + + output, err := cmd.CombinedOutput() + if err != nil { + return fmt.Errorf("git submodule init failed: %s - %w", strings.TrimSpace(string(output)), err) + } + return nil +} + +func (c *Cloner) buildEnv() []string { + env := os.Environ() + + if c.config.SSHKeyPath != "" { + sshCmd := fmt.Sprintf("ssh -i %s -o StrictHostKeyChecking=no", c.config.SSHKeyPath) + env = append(env, "GIT_SSH_COMMAND="+sshCmd) + } + + return env +} + +func (c *Cloner) IsConfigured() bool { + return c.config.URL != "" +} diff --git a/pkg/harness/git/manager.go b/pkg/harness/git/manager.go new file mode 100644 index 00000000..802560e0 --- /dev/null +++ b/pkg/harness/git/manager.go @@ -0,0 +1,222 @@ +package git + +import ( + "context" + "fmt" + "os/exec" + "regexp" + "strings" + "time" +) + +type Manager struct { + workspacePath string + branchPrefix string + currentBranch string + originalBranch string +} + +type ManagerOption func(*Manager) + +func WithBranchPrefix(prefix string) ManagerOption { + return func(m *Manager) { + m.branchPrefix = prefix + } +} + +func NewManager(workspacePath string, opts ...ManagerOption) *Manager { + m := &Manager{ + workspacePath: workspacePath, + branchPrefix: "agent/", + } + + for _, opt := range opts { + opt(m) + } + + return m +} + +func (m *Manager) CreateBranch(ctx context.Context, task string, agentID string) (string, error) { + original, err := m.GetCurrentBranch(ctx) + if err != nil { + return "", fmt.Errorf("failed to get current branch: %w", err) + } + m.originalBranch = original + + slug := slugify(task) + if len(slug) > 30 { + slug = slug[:30] + } + + timestamp := time.Now().Format("20060102-150405") + branchName := fmt.Sprintf("%s%s-%s-%s", m.branchPrefix, slug, timestamp, agentID) + + cmd := exec.CommandContext(ctx, "git", "checkout", "-b", branchName) + cmd.Dir = m.workspacePath + + output, err := cmd.CombinedOutput() + if err != nil { + return "", fmt.Errorf("failed to create branch: %s - %w", string(output), err) + } + + m.currentBranch = branchName + return branchName, nil +} + +func (m *Manager) SwitchToBranch(ctx context.Context, branchName string) error { + cmd := exec.CommandContext(ctx, "git", "checkout", branchName) + cmd.Dir = m.workspacePath + + output, err := cmd.CombinedOutput() + if err != nil { + return fmt.Errorf("failed to switch branch: %s - %w", string(output), err) + } + + m.currentBranch = branchName + return nil +} + +func (m *Manager) Commit(ctx context.Context, message string) (string, error) { + addCmd := exec.CommandContext(ctx, "git", "add", "-A") + addCmd.Dir = m.workspacePath + + if output, err := addCmd.CombinedOutput(); err != nil { + return "", fmt.Errorf("failed to stage changes: %s - %w", string(output), err) + } + + statusCmd := exec.CommandContext(ctx, "git", "status", "--porcelain") + statusCmd.Dir = m.workspacePath + + statusOutput, err := statusCmd.CombinedOutput() + if err != nil { + return "", fmt.Errorf("failed to check status: %w", err) + } + + if len(strings.TrimSpace(string(statusOutput))) == 0 { + return "", nil + } + + commitCmd := exec.CommandContext(ctx, "git", "commit", "-m", message) + commitCmd.Dir = m.workspacePath + + if output, err := commitCmd.CombinedOutput(); err != nil { + return "", fmt.Errorf("failed to commit: %s - %w", string(output), err) + } + + shaCmd := exec.CommandContext(ctx, "git", "rev-parse", "HEAD") + shaCmd.Dir = m.workspacePath + + shaOutput, err := shaCmd.CombinedOutput() + if err != nil { + return "", fmt.Errorf("failed to get commit SHA: %w", err) + } + + return strings.TrimSpace(string(shaOutput)), nil +} + +func (m *Manager) Push(ctx context.Context) error { + cmd := exec.CommandContext(ctx, "git", "push", "-u", "origin", m.currentBranch) + cmd.Dir = m.workspacePath + + output, err := cmd.CombinedOutput() + if err != nil { + return fmt.Errorf("failed to push: %s - %w", string(output), err) + } + + return nil +} + +func (m *Manager) GetCurrentBranch(ctx context.Context) (string, error) { + cmd := exec.CommandContext(ctx, "git", "branch", "--show-current") + cmd.Dir = m.workspacePath + + output, err := cmd.CombinedOutput() + if err != nil { + return "", fmt.Errorf("failed to get current branch: %s - %w", string(output), err) + } + + return strings.TrimSpace(string(output)), nil +} + +func (m *Manager) GetStatus(ctx context.Context) (string, error) { + cmd := exec.CommandContext(ctx, "git", "status", "--short") + cmd.Dir = m.workspacePath + + output, err := cmd.CombinedOutput() + if err != nil { + return "", fmt.Errorf("failed to get status: %s - %w", string(output), err) + } + + return string(output), nil +} + +func (m *Manager) GetDiff(ctx context.Context) (string, error) { + cmd := exec.CommandContext(ctx, "git", "diff", "HEAD") + cmd.Dir = m.workspacePath + + output, err := cmd.CombinedOutput() + if err != nil { + return "", fmt.Errorf("failed to get diff: %s - %w", string(output), err) + } + + return string(output), nil +} + +func (m *Manager) RestoreOriginalBranch(ctx context.Context) error { + if m.originalBranch == "" { + return nil + } + + return m.SwitchToBranch(ctx, m.originalBranch) +} + +func (m *Manager) HasUncommittedChanges(ctx context.Context) (bool, error) { + status, err := m.GetStatus(ctx) + if err != nil { + return false, err + } + + return len(strings.TrimSpace(status)) > 0, nil +} + +func (m *Manager) IsGitRepo(ctx context.Context) bool { + cmd := exec.CommandContext(ctx, "git", "rev-parse", "--git-dir") + cmd.Dir = m.workspacePath + + return cmd.Run() == nil +} + +func (m *Manager) GetRecentCommits(ctx context.Context, count int) ([]string, error) { + cmd := exec.CommandContext(ctx, "git", "log", fmt.Sprintf("-%d", count), "--oneline") + cmd.Dir = m.workspacePath + + output, err := cmd.CombinedOutput() + if err != nil { + return nil, fmt.Errorf("failed to get commits: %s - %w", string(output), err) + } + + lines := strings.Split(strings.TrimSpace(string(output)), "\n") + if len(lines) == 1 && lines[0] == "" { + return []string{}, nil + } + + return lines, nil +} + +func slugify(s string) string { + s = strings.ToLower(s) + + reg := regexp.MustCompile(`[^a-z0-9\s-]`) + s = reg.ReplaceAllString(s, "") + + reg = regexp.MustCompile(`[\s_]+`) + s = reg.ReplaceAllString(s, "-") + + reg = regexp.MustCompile(`-+`) + s = reg.ReplaceAllString(s, "-") + + s = strings.Trim(s, "-") + + return s +} diff --git a/pkg/harness/git/manager_test.go b/pkg/harness/git/manager_test.go new file mode 100644 index 00000000..deb7b82a --- /dev/null +++ b/pkg/harness/git/manager_test.go @@ -0,0 +1,249 @@ +package git + +import ( + "context" + "os" + "os/exec" + "path/filepath" + "strings" + "testing" +) + +func setupGitRepo(t *testing.T) string { + tmpDir := t.TempDir() + + cmd := exec.Command("git", "init") + cmd.Dir = tmpDir + if err := cmd.Run(); err != nil { + t.Fatalf("failed to init git repo: %v", err) + } + + cmd = exec.Command("git", "config", "user.email", "test@test.com") + cmd.Dir = tmpDir + cmd.Run() + + cmd = exec.Command("git", "config", "user.name", "Test User") + cmd.Dir = tmpDir + cmd.Run() + + os.WriteFile(filepath.Join(tmpDir, "README.md"), []byte("# Test"), 0644) + + cmd = exec.Command("git", "add", ".") + cmd.Dir = tmpDir + cmd.Run() + + cmd = exec.Command("git", "commit", "-m", "Initial commit") + cmd.Dir = tmpDir + cmd.Run() + + return tmpDir +} + +func TestManager_IsGitRepo(t *testing.T) { + tmpDir := setupGitRepo(t) + ctx := context.Background() + + m := NewManager(tmpDir) + + if !m.IsGitRepo(ctx) { + t.Error("should detect git repo") + } + + nonGitDir := t.TempDir() + m2 := NewManager(nonGitDir) + + if m2.IsGitRepo(ctx) { + t.Error("should not detect git repo in non-git directory") + } +} + +func TestManager_GetCurrentBranch(t *testing.T) { + tmpDir := setupGitRepo(t) + ctx := context.Background() + + m := NewManager(tmpDir) + + branch, err := m.GetCurrentBranch(ctx) + if err != nil { + t.Fatalf("GetCurrentBranch failed: %v", err) + } + + if branch != "main" && branch != "master" { + t.Errorf("GetCurrentBranch returned %q, want main or master", branch) + } +} + +func TestManager_CreateBranch(t *testing.T) { + tmpDir := setupGitRepo(t) + ctx := context.Background() + + m := NewManager(tmpDir, WithBranchPrefix("test/")) + + branchName, err := m.CreateBranch(ctx, "Fix the bug", "agent-1") + if err != nil { + t.Fatalf("CreateBranch failed: %v", err) + } + + if !strings.HasPrefix(branchName, "test/fix-the-bug-") { + t.Errorf("branch name %q doesn't have expected prefix", branchName) + } + + if !strings.HasSuffix(branchName, "-agent-1") { + t.Errorf("branch name %q doesn't have expected suffix", branchName) + } + + currentBranch, _ := m.GetCurrentBranch(ctx) + if currentBranch != branchName { + t.Errorf("current branch is %q, want %q", currentBranch, branchName) + } +} + +func TestManager_Commit(t *testing.T) { + tmpDir := setupGitRepo(t) + ctx := context.Background() + + m := NewManager(tmpDir) + + os.WriteFile(filepath.Join(tmpDir, "newfile.txt"), []byte("new content"), 0644) + + sha, err := m.Commit(ctx, "Add new file") + if err != nil { + t.Fatalf("Commit failed: %v", err) + } + + if sha == "" { + t.Error("commit should return SHA") + } + + if len(sha) != 40 { + t.Errorf("SHA length = %d, want 40", len(sha)) + } +} + +func TestManager_CommitNoChanges(t *testing.T) { + tmpDir := setupGitRepo(t) + ctx := context.Background() + + m := NewManager(tmpDir) + + sha, err := m.Commit(ctx, "No changes") + if err != nil { + t.Fatalf("Commit failed: %v", err) + } + + if sha != "" { + t.Error("commit with no changes should return empty SHA") + } +} + +func TestManager_HasUncommittedChanges(t *testing.T) { + tmpDir := setupGitRepo(t) + ctx := context.Background() + + m := NewManager(tmpDir) + + hasChanges, err := m.HasUncommittedChanges(ctx) + if err != nil { + t.Fatalf("HasUncommittedChanges failed: %v", err) + } + if hasChanges { + t.Error("should have no uncommitted changes initially") + } + + os.WriteFile(filepath.Join(tmpDir, "modified.txt"), []byte("content"), 0644) + + hasChanges, err = m.HasUncommittedChanges(ctx) + if err != nil { + t.Fatalf("HasUncommittedChanges failed: %v", err) + } + if !hasChanges { + t.Error("should have uncommitted changes after creating file") + } +} + +func TestManager_GetStatus(t *testing.T) { + tmpDir := setupGitRepo(t) + ctx := context.Background() + + m := NewManager(tmpDir) + + os.WriteFile(filepath.Join(tmpDir, "newfile.txt"), []byte("content"), 0644) + + status, err := m.GetStatus(ctx) + if err != nil { + t.Fatalf("GetStatus failed: %v", err) + } + + if !strings.Contains(status, "newfile.txt") { + t.Errorf("status should contain newfile.txt: %q", status) + } +} + +func TestManager_GetRecentCommits(t *testing.T) { + tmpDir := setupGitRepo(t) + ctx := context.Background() + + m := NewManager(tmpDir) + + commits, err := m.GetRecentCommits(ctx, 5) + if err != nil { + t.Fatalf("GetRecentCommits failed: %v", err) + } + + if len(commits) != 1 { + t.Errorf("expected 1 commit, got %d", len(commits)) + } + + if !strings.Contains(commits[0], "Initial commit") { + t.Errorf("commit message should contain 'Initial commit': %q", commits[0]) + } +} + +func TestSlugify(t *testing.T) { + tests := []struct { + input string + want string + }{ + {"Fix the bug", "fix-the-bug"}, + {"Add new feature!", "add-new-feature"}, + {"Update README.md", "update-readmemd"}, + {" spaces around ", "spaces-around"}, + {"UPPERCASE text", "uppercase-text"}, + {"multiple---dashes", "multiple-dashes"}, + {"special@#$chars", "specialchars"}, + } + + for _, tt := range tests { + t.Run(tt.input, func(t *testing.T) { + got := slugify(tt.input) + if got != tt.want { + t.Errorf("slugify(%q) = %q, want %q", tt.input, got, tt.want) + } + }) + } +} + +func TestManager_RestoreOriginalBranch(t *testing.T) { + tmpDir := setupGitRepo(t) + ctx := context.Background() + + m := NewManager(tmpDir) + + originalBranch, _ := m.GetCurrentBranch(ctx) + + m.CreateBranch(ctx, "Feature branch", "agent-1") + + currentBranch, _ := m.GetCurrentBranch(ctx) + if currentBranch == originalBranch { + t.Error("should be on new branch") + } + + if err := m.RestoreOriginalBranch(ctx); err != nil { + t.Fatalf("RestoreOriginalBranch failed: %v", err) + } + + currentBranch, _ = m.GetCurrentBranch(ctx) + if currentBranch != originalBranch { + t.Errorf("should be back on original branch %q, got %q", originalBranch, currentBranch) + } +} diff --git a/pkg/harness/hooks.go b/pkg/harness/hooks.go new file mode 100644 index 00000000..b7312b08 --- /dev/null +++ b/pkg/harness/hooks.go @@ -0,0 +1,282 @@ +// Package harness provides an agentic execution harness for Station. +package harness + +import ( + "context" + "fmt" + "strings" + "sync" + + "github.com/firebase/genkit/go/ai" +) + +// HookResult indicates how to proceed after a hook runs. +type HookResult string + +const ( + // HookContinue allows the operation to proceed normally. + HookContinue HookResult = "continue" + // HookBlock stops the operation with an error returned to the LLM. + HookBlock HookResult = "block" + // HookInterrupt pauses the operation for human approval. + HookInterrupt HookResult = "interrupt" +) + +// PreHookFunc is called before a tool executes. +// It receives the tool request and returns a result indicating how to proceed. +type PreHookFunc func(ctx context.Context, req *ai.ToolRequest) (HookResult, string) + +// PostHookFunc is called after a tool executes successfully. +// It receives the tool request and output for logging/monitoring purposes. +type PostHookFunc func(ctx context.Context, req *ai.ToolRequest, output interface{}) + +// HookRegistry manages pre and post tool execution hooks. +type HookRegistry struct { + mu sync.RWMutex + preHooks []PreHookFunc + postHooks []PostHookFunc +} + +// NewHookRegistry creates a new hook registry. +func NewHookRegistry() *HookRegistry { + return &HookRegistry{ + preHooks: make([]PreHookFunc, 0), + postHooks: make([]PostHookFunc, 0), + } +} + +// RegisterPreHook adds a pre-execution hook. +func (hr *HookRegistry) RegisterPreHook(hook PreHookFunc) { + hr.mu.Lock() + defer hr.mu.Unlock() + hr.preHooks = append(hr.preHooks, hook) +} + +// RegisterPostHook adds a post-execution hook. +func (hr *HookRegistry) RegisterPostHook(hook PostHookFunc) { + hr.mu.Lock() + defer hr.mu.Unlock() + hr.postHooks = append(hr.postHooks, hook) +} + +// RunPreHooks executes all pre-hooks and returns the most restrictive result. +// Priority: Interrupt > Block > Continue +func (hr *HookRegistry) RunPreHooks(ctx context.Context, req *ai.ToolRequest) (HookResult, string) { + hr.mu.RLock() + defer hr.mu.RUnlock() + + var finalResult HookResult = HookContinue + var finalMsg string + + for _, hook := range hr.preHooks { + result, msg := hook(ctx, req) + switch result { + case HookInterrupt: + // Interrupt has highest priority - return immediately + return HookInterrupt, msg + case HookBlock: + // Block has higher priority than Continue + finalResult = HookBlock + finalMsg = msg + case HookContinue: + // Continue only if nothing else is blocking + if finalResult == HookContinue { + finalResult = HookContinue + } + } + } + + return finalResult, finalMsg +} + +// RunPostHooks executes all post-hooks. +func (hr *HookRegistry) RunPostHooks(ctx context.Context, req *ai.ToolRequest, output interface{}) { + hr.mu.RLock() + defer hr.mu.RUnlock() + + for _, hook := range hr.postHooks { + hook(ctx, req, output) + } +} + +// RegisterDoomLoopDetection registers a doom loop detection hook. +func (hr *HookRegistry) RegisterDoomLoopDetection(detector *DoomLoopDetector) { + hr.RegisterPreHook(func(ctx context.Context, req *ai.ToolRequest) (HookResult, string) { + if detector.IsInDoomLoop(req.Name, req.Input) { + return HookInterrupt, fmt.Sprintf( + "Doom loop detected: tool '%s' has been called %d times consecutively with similar arguments. "+ + "This may indicate the agent is stuck. Consider rephrasing the task or checking the tool output.", + req.Name, + detector.threshold, + ) + } + return HookContinue, "" + }) +} + +// RegisterPermissionCheck registers permission checking hooks based on config. +func (hr *HookRegistry) RegisterPermissionCheck(perms PermissionsConfig) { + hr.RegisterPreHook(func(ctx context.Context, req *ai.ToolRequest) (HookResult, string) { + // Handle bash tool permission checks + if req.Name == "bash" || req.Name == "execute_command" { + command := extractCommand(req.Input) + if command != "" { + return checkBashPermission(command, perms.Bash) + } + } + return HookContinue, "" + }) +} + +// extractCommand extracts the command string from tool input. +func extractCommand(input interface{}) string { + switch v := input.(type) { + case map[string]interface{}: + if cmd, ok := v["command"].(string); ok { + return cmd + } + case string: + return v + } + return "" +} + +// checkBashPermission checks a command against permission rules. +// Rules are matched in order; first match wins. +func checkBashPermission(command string, rules map[string]string) (HookResult, string) { + // Check specific rules first (longer patterns) + type ruleScore struct { + pattern string + action string + score int + } + + var matches []ruleScore + for pattern, action := range rules { + if matchWildcard(pattern, command) { + // Score by specificity (longer patterns are more specific) + matches = append(matches, ruleScore{pattern, action, len(pattern)}) + } + } + + // Sort by specificity (most specific first) + // Use the most specific matching rule + var bestMatch *ruleScore + for i := range matches { + if bestMatch == nil || matches[i].score > bestMatch.score { + bestMatch = &matches[i] + } + } + + if bestMatch != nil { + switch bestMatch.action { + case "deny": + return HookBlock, fmt.Sprintf("command '%s' is denied by rule '%s'", command, bestMatch.pattern) + case "ask": + return HookInterrupt, fmt.Sprintf("command '%s' requires approval (matched rule '%s')", command, bestMatch.pattern) + case "allow": + return HookContinue, "" + } + } + + // Default: allow if no rules match + return HookContinue, "" +} + +// matchWildcard performs OpenCode-style wildcard matching. +// "*" matches any sequence of characters. +func matchWildcard(pattern, str string) bool { + // Handle empty pattern + if pattern == "" { + return str == "" + } + + // Split pattern by wildcards + parts := strings.Split(pattern, "*") + + // If pattern is just "*", it matches everything + if len(parts) == 1 { + return pattern == str + } + + // Check if string starts with first part (if pattern doesn't start with *) + pos := 0 + if parts[0] != "" { + if !strings.HasPrefix(str, parts[0]) { + return false + } + pos = len(parts[0]) + } + + // Check middle parts + for i := 1; i < len(parts)-1; i++ { + if parts[i] == "" { + continue + } + idx := strings.Index(str[pos:], parts[i]) + if idx == -1 { + return false + } + pos += idx + len(parts[i]) + } + + // Check if string ends with last part (if pattern doesn't end with *) + lastPart := parts[len(parts)-1] + if lastPart != "" { + return strings.HasSuffix(str, lastPart) && pos <= len(str)-len(lastPart) + } + + return true +} + +// PathPermissionHook creates a hook that checks file path permissions. +func PathPermissionHook(workspacePath string, externalPolicy string) PreHookFunc { + return func(ctx context.Context, req *ai.ToolRequest) (HookResult, string) { + // File operation tools + fileTools := map[string]bool{ + "read_file": true, + "write_file": true, + "edit_file": true, + "list_files": true, + "glob": true, + "grep": true, + } + + if !fileTools[req.Name] { + return HookContinue, "" + } + + path := extractPath(req.Input) + if path == "" { + return HookContinue, "" + } + + // Check if path is outside workspace + if !strings.HasPrefix(path, workspacePath) { + switch externalPolicy { + case "deny": + return HookBlock, fmt.Sprintf("access to path '%s' outside workspace is denied", path) + case "ask": + return HookInterrupt, fmt.Sprintf("access to path '%s' outside workspace requires approval", path) + } + } + + return HookContinue, "" + } +} + +// extractPath extracts the file path from tool input. +func extractPath(input interface{}) string { + switch v := input.(type) { + case map[string]interface{}: + // Try common path parameter names + for _, key := range []string{"path", "file_path", "filePath", "file", "directory"} { + if p, ok := v[key].(string); ok { + return p + } + } + case string: + return v + } + return "" +} diff --git a/pkg/harness/hooks_test.go b/pkg/harness/hooks_test.go new file mode 100644 index 00000000..4d78f2ae --- /dev/null +++ b/pkg/harness/hooks_test.go @@ -0,0 +1,170 @@ +package harness + +import ( + "context" + "testing" + + "github.com/firebase/genkit/go/ai" +) + +func TestHookResult_Constants(t *testing.T) { + if HookContinue != "continue" { + t.Errorf("HookContinue = %q, want %q", HookContinue, "continue") + } + if HookBlock != "block" { + t.Errorf("HookBlock = %q, want %q", HookBlock, "block") + } + if HookInterrupt != "interrupt" { + t.Errorf("HookInterrupt = %q, want %q", HookInterrupt, "interrupt") + } +} + +func TestHookRegistry_PreHooks(t *testing.T) { + registry := NewHookRegistry() + + hookCalled := false + registry.RegisterPreHook(func(ctx context.Context, req *ai.ToolRequest) (HookResult, string) { + hookCalled = true + return HookContinue, "" + }) + + req := &ai.ToolRequest{Name: "test_tool"} + result, _ := registry.RunPreHooks(context.Background(), req) + + if !hookCalled { + t.Error("PreHook was not called") + } + if result != HookContinue { + t.Errorf("result = %v, want %v", result, HookContinue) + } +} + +func TestHookRegistry_PreHooks_Priority(t *testing.T) { + registry := NewHookRegistry() + + registry.RegisterPreHook(func(ctx context.Context, req *ai.ToolRequest) (HookResult, string) { + return HookContinue, "" + }) + + registry.RegisterPreHook(func(ctx context.Context, req *ai.ToolRequest) (HookResult, string) { + return HookBlock, "blocked" + }) + + req := &ai.ToolRequest{Name: "test_tool"} + result, msg := registry.RunPreHooks(context.Background(), req) + + if result != HookBlock { + t.Errorf("result = %v, want %v", result, HookBlock) + } + if msg != "blocked" { + t.Errorf("msg = %q, want %q", msg, "blocked") + } +} + +func TestHookRegistry_PreHooks_InterruptPriority(t *testing.T) { + registry := NewHookRegistry() + + registry.RegisterPreHook(func(ctx context.Context, req *ai.ToolRequest) (HookResult, string) { + return HookBlock, "blocked" + }) + + registry.RegisterPreHook(func(ctx context.Context, req *ai.ToolRequest) (HookResult, string) { + return HookInterrupt, "needs approval" + }) + + req := &ai.ToolRequest{Name: "test_tool"} + result, msg := registry.RunPreHooks(context.Background(), req) + + if result != HookInterrupt { + t.Errorf("result = %v, want %v (Interrupt has highest priority)", result, HookInterrupt) + } + if msg != "needs approval" { + t.Errorf("msg = %q, want %q", msg, "needs approval") + } +} + +func TestHookRegistry_PostHooks(t *testing.T) { + registry := NewHookRegistry() + + hookCalled := false + var capturedOutput interface{} + + registry.RegisterPostHook(func(ctx context.Context, req *ai.ToolRequest, output interface{}) { + hookCalled = true + capturedOutput = output + }) + + req := &ai.ToolRequest{Name: "test_tool"} + registry.RunPostHooks(context.Background(), req, "test output") + + if !hookCalled { + t.Error("PostHook was not called") + } + if capturedOutput != "test output" { + t.Errorf("capturedOutput = %v, want %v", capturedOutput, "test output") + } +} + +func TestMatchWildcard(t *testing.T) { + tests := []struct { + pattern string + str string + want bool + }{ + {"*", "anything", true}, + {"*", "", true}, + {"rm -rf *", "rm -rf /", true}, + {"rm -rf *", "rm -rf /home/user", true}, + {"rm -rf *", "rm /home", false}, + {"git push *", "git push origin main", true}, + {"git push *", "git pull", false}, + {"git push --force*", "git push --force origin main", true}, + {"git push --force*", "git push origin main", false}, + {"*.txt", "file.txt", true}, + {"*.txt", "file.log", false}, + {"test*end", "test123end", true}, + {"test*end", "testend", true}, + {"test*end", "test123", false}, + {"exact", "exact", true}, + {"exact", "exactnot", false}, + } + + for _, tt := range tests { + t.Run(tt.pattern+"_"+tt.str, func(t *testing.T) { + got := matchWildcard(tt.pattern, tt.str) + if got != tt.want { + t.Errorf("matchWildcard(%q, %q) = %v, want %v", tt.pattern, tt.str, got, tt.want) + } + }) + } +} + +func TestCheckBashPermission(t *testing.T) { + rules := map[string]string{ + "*": "allow", + "rm -rf *": "deny", + "git push *": "ask", + "git push --force*": "deny", + } + + tests := []struct { + command string + wantResult HookResult + }{ + {"ls -la", HookContinue}, + {"cat file.txt", HookContinue}, + {"rm -rf /", HookBlock}, + {"rm -rf /home/user", HookBlock}, + {"git push origin main", HookInterrupt}, + {"git push --force origin main", HookBlock}, + } + + for _, tt := range tests { + t.Run(tt.command, func(t *testing.T) { + result, _ := checkBashPermission(tt.command, rules) + if result != tt.wantResult { + t.Errorf("checkBashPermission(%q) = %v, want %v", tt.command, result, tt.wantResult) + } + }) + } +} diff --git a/pkg/harness/memory/fs_backend.go b/pkg/harness/memory/fs_backend.go new file mode 100644 index 00000000..2b96b4f8 --- /dev/null +++ b/pkg/harness/memory/fs_backend.go @@ -0,0 +1,19 @@ +package memory + +import ( + "os" + "path/filepath" +) + +// FSBackend implements Backend using the local filesystem +type FSBackend struct{} + +// NewFSBackend creates a new filesystem backend +func NewFSBackend() *FSBackend { + return &FSBackend{} +} + +// ReadFile reads a file's content +func (b *FSBackend) ReadFile(path string) ([]byte, error) { + return os.ReadFile(filepath.Clean(path)) +} diff --git a/pkg/harness/memory/middleware.go b/pkg/harness/memory/middleware.go new file mode 100644 index 00000000..6e58dd1d --- /dev/null +++ b/pkg/harness/memory/middleware.go @@ -0,0 +1,116 @@ +package memory + +import ( + "fmt" + "log/slog" + "os" + "path/filepath" + "strings" +) + +// MemoryMiddleware loads and injects AGENTS.md content into the system prompt +type MemoryMiddleware struct { + backend Backend + sources []string + logger *slog.Logger +} + +// NewMemoryMiddleware creates a new memory middleware with the given backend and sources +func NewMemoryMiddleware(backend Backend, sources []string) *MemoryMiddleware { + return &MemoryMiddleware{ + backend: backend, + sources: sources, + logger: slog.Default().With("component", "memory_middleware"), + } +} + +// LoadMemory reads all AGENTS.md files from configured sources +// Returns a map of source path to content +func (m *MemoryMiddleware) LoadMemory() (map[string]string, error) { + contents := make(map[string]string) + + for _, source := range m.sources { + expandedPath := expandPath(source) + content, err := m.backend.ReadFile(expandedPath) + if err != nil { + // Memory files are optional - just log and continue + m.logger.Debug("memory source not found", + "source", expandedPath, + "error", err) + continue + } + + contents[source] = string(content) + m.logger.Debug("loaded memory", + "source", expandedPath, + "bytes", len(content)) + } + + return contents, nil +} + +// FormatSystemPromptSection formats memory content for injection into system prompt +// Memory is always loaded (unlike skills which use progressive disclosure) +func (m *MemoryMiddleware) FormatSystemPromptSection() (string, error) { + contents, err := m.LoadMemory() + if err != nil { + return "", err + } + + if len(contents) == 0 { + return "", nil + } + + var sb strings.Builder + sb.WriteString("\n\n\n") + + for source, content := range contents { + sb.WriteString(fmt.Sprintf("\n\n", source)) + sb.WriteString(content) + sb.WriteString("\n") + } + + sb.WriteString("\n") + sb.WriteString(` + +The above context was loaded from AGENTS.md files in your filesystem. As you learn from interactions: + +**When to update memory (use edit_file on AGENTS.md):** +- User explicitly asks you to remember something +- User describes role or behavior preferences +- User provides feedback on your work that should persist +- You discover new patterns or preferences + +**When NOT to update memory:** +- Temporary or transient information +- One-time task requests +- Simple questions that don't reveal lasting preferences + +**Security:** +- NEVER store API keys, passwords, or credentials in memory files + +`) + + return sb.String(), nil +} + +// expandPath expands ~ to the user's home directory +func expandPath(path string) string { + if strings.HasPrefix(path, "~/") { + home, err := os.UserHomeDir() + if err != nil { + return path + } + return filepath.Join(home, path[2:]) + } + return path +} + +// DefaultMemorySources returns the default memory source paths +func DefaultMemorySources(envPath string) []string { + return []string{ + "~/.config/station/AGENTS.md", + filepath.Join(envPath, "memory", "AGENTS.md"), + ".station/AGENTS.md", + } +} diff --git a/pkg/harness/memory/middleware_test.go b/pkg/harness/memory/middleware_test.go new file mode 100644 index 00000000..54690fc6 --- /dev/null +++ b/pkg/harness/memory/middleware_test.go @@ -0,0 +1,188 @@ +package memory + +import ( + "strings" + "testing" +) + +// MockBackend implements Backend for testing +type MockBackend struct { + Files map[string][]byte +} + +func NewMockBackend() *MockBackend { + return &MockBackend{ + Files: make(map[string][]byte), + } +} + +func (m *MockBackend) ReadFile(path string) ([]byte, error) { + if content, ok := m.Files[path]; ok { + return content, nil + } + return nil, &mockError{msg: "file not found: " + path} +} + +type mockError struct { + msg string +} + +func (e *mockError) Error() string { + return e.msg +} + +func TestMemoryMiddleware_LoadMemory(t *testing.T) { + backend := NewMockBackend() + + backend.Files["/global/AGENTS.md"] = []byte(`# Global Memory +This is global context.`) + + backend.Files["/project/AGENTS.md"] = []byte(`# Project Memory +This is project-specific context.`) + + middleware := NewMemoryMiddleware(backend, []string{ + "/global/AGENTS.md", + "/project/AGENTS.md", + }) + + contents, err := middleware.LoadMemory() + if err != nil { + t.Fatalf("LoadMemory() error = %v", err) + } + + if len(contents) != 2 { + t.Errorf("LoadMemory() returned %d contents, want 2", len(contents)) + } + + if !strings.Contains(contents["/global/AGENTS.md"], "Global Memory") { + t.Error("missing global memory content") + } + + if !strings.Contains(contents["/project/AGENTS.md"], "Project Memory") { + t.Error("missing project memory content") + } +} + +func TestMemoryMiddleware_FormatSystemPromptSection(t *testing.T) { + backend := NewMockBackend() + + backend.Files["/memory/AGENTS.md"] = []byte(`# Build Commands +- make build +- make test + +# Code Style +Use slog for logging`) + + middleware := NewMemoryMiddleware(backend, []string{"/memory/AGENTS.md"}) + + section, err := middleware.FormatSystemPromptSection() + if err != nil { + t.Fatalf("FormatSystemPromptSection() error = %v", err) + } + + // Check key elements + if !strings.Contains(section, "") { + t.Error("section should contain agent_memory tag") + } + if !strings.Contains(section, "") { + t.Error("section should contain closing agent_memory tag") + } + if !strings.Contains(section, "Build Commands") { + t.Error("section should contain memory content") + } + if !strings.Contains(section, "Source: /memory/AGENTS.md") { + t.Error("section should contain source path comment") + } + if !strings.Contains(section, "memory_guidelines") { + t.Error("section should contain guidelines") + } + if !strings.Contains(section, "NEVER store API keys") { + t.Error("section should contain security warning") + } +} + +func TestMemoryMiddleware_EmptyMemory(t *testing.T) { + backend := NewMockBackend() + // No files set up + + middleware := NewMemoryMiddleware(backend, []string{"/nonexistent/AGENTS.md"}) + + section, err := middleware.FormatSystemPromptSection() + if err != nil { + t.Fatalf("FormatSystemPromptSection() should not error, got %v", err) + } + + if section != "" { + t.Errorf("empty memory should produce empty section, got %q", section) + } +} + +func TestMemoryMiddleware_MissingSource(t *testing.T) { + backend := NewMockBackend() + // Only one file exists + backend.Files["/exists/AGENTS.md"] = []byte("# Exists") + + middleware := NewMemoryMiddleware(backend, []string{ + "/nonexistent/AGENTS.md", + "/exists/AGENTS.md", + }) + + contents, err := middleware.LoadMemory() + if err != nil { + t.Fatalf("LoadMemory() should not error, got %v", err) + } + + if len(contents) != 1 { + t.Errorf("LoadMemory() returned %d contents, want 1", len(contents)) + } + + if _, ok := contents["/exists/AGENTS.md"]; !ok { + t.Error("should have loaded existing file") + } +} + +func TestExpandPath(t *testing.T) { + tests := []struct { + input string + want string + }{ + {"/absolute/path", "/absolute/path"}, + {"relative/path", "relative/path"}, + // Note: can't easily test ~ expansion without mocking os.UserHomeDir + } + + for _, tt := range tests { + t.Run(tt.input, func(t *testing.T) { + got := expandPath(tt.input) + if !strings.HasPrefix(tt.input, "~/") && got != tt.want { + t.Errorf("expandPath(%q) = %q, want %q", tt.input, got, tt.want) + } + }) + } +} + +func TestDefaultMemorySources(t *testing.T) { + sources := DefaultMemorySources("/env/default") + + if len(sources) < 2 { + t.Error("should have at least 2 default sources") + } + + hasGlobal := false + hasProject := false + for _, s := range sources { + if strings.Contains(s, ".config/station") { + hasGlobal = true + } + if strings.Contains(s, ".station") { + hasProject = true + } + } + + if !hasGlobal { + t.Error("should have global config source") + } + if !hasProject { + t.Error("should have project source") + } +} diff --git a/pkg/harness/memory/types.go b/pkg/harness/memory/types.go new file mode 100644 index 00000000..6f2a1ed1 --- /dev/null +++ b/pkg/harness/memory/types.go @@ -0,0 +1,11 @@ +package memory + +// MemoryConfig configures the memory middleware from agent frontmatter +type MemoryConfig struct { + Sources []string `yaml:"sources,omitempty" json:"sources,omitempty"` +} + +// Backend interface for file operations +type Backend interface { + ReadFile(path string) ([]byte, error) +} diff --git a/pkg/harness/multi_agent_e2e_test.go b/pkg/harness/multi_agent_e2e_test.go new file mode 100644 index 00000000..21a6556a --- /dev/null +++ b/pkg/harness/multi_agent_e2e_test.go @@ -0,0 +1,491 @@ +//go:build integration + +package harness_test + +import ( + "bytes" + "context" + "fmt" + "io" + "net" + "os" + "strings" + "testing" + "time" + + "station/internal/config" + "station/internal/genkit/anthropic_oauth" + "station/internal/lattice" + "station/internal/logging" + "station/pkg/harness" + harnessNats "station/pkg/harness/nats" + "station/pkg/harness/tools" + "station/pkg/harness/workspace" + + "github.com/firebase/genkit/go/genkit" + "github.com/firebase/genkit/go/plugins/compat_oai/openai" + "github.com/firebase/genkit/go/plugins/googlegenai" +) + +func getFreePortForE2E(t *testing.T) int { + t.Helper() + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("Failed to find free port: %v", err) + } + port := listener.Addr().(*net.TCPAddr).Port + listener.Close() + return port +} + +func TestMultiAgentWorkflow_E2E_RealLLM_RealNATS(t *testing.T) { + if os.Getenv("HARNESS_E2E_TEST") != "1" { + t.Skip("Skipping e2e test. Set HARNESS_E2E_TEST=1 to run with real LLM + real NATS") + } + + ctx := context.Background() + + t.Log("=== MULTI-AGENT WORKFLOW E2E TEST ===") + t.Log("This test runs two AI agents that collaborate via NATS:") + t.Log(" Agent A (Analyzer): Analyzes buggy code, stores findings in NATS") + t.Log(" Agent B (Fixer): Reads findings from NATS, fixes the code") + t.Log("") + + if err := config.InitViper(""); err != nil { + t.Fatalf("Failed to init viper: %v", err) + } + + cfg, err := config.Load() + if err != nil { + t.Fatalf("Failed to load config: %v", err) + } + + t.Logf("Using AI: %s/%s", cfg.AIProvider, cfg.AIModel) + + genkitApp, err := initGenkitForMultiAgent(ctx, cfg) + if err != nil { + t.Fatalf("Failed to initialize GenKit: %v", err) + } + + port := getFreePortForE2E(t) + httpPort := getFreePortForE2E(t) + + serverCfg := config.LatticeEmbeddedNATSConfig{ + Port: port, + HTTPPort: httpPort, + StoreDir: t.TempDir(), + } + + t.Log("Starting embedded NATS server with JetStream...") + server := lattice.NewEmbeddedServer(serverCfg) + if err := server.Start(); err != nil { + t.Fatalf("Failed to start embedded NATS: %v", err) + } + defer server.Shutdown() + t.Logf("NATS server running at %s", server.ClientURL()) + + clientCfg := config.LatticeConfig{ + StationID: "multi-agent-test", + StationName: "Multi-Agent E2E Test", + NATS: config.LatticeNATSConfig{URL: server.ClientURL()}, + } + + client, err := lattice.NewClient(clientCfg) + if err != nil { + t.Fatalf("Failed to create lattice client: %v", err) + } + if err := client.Connect(); err != nil { + t.Fatalf("Failed to connect to NATS: %v", err) + } + defer client.Close() + + natsStore, err := harnessNats.NewStore(client.Conn(), harnessNats.DefaultStoreConfig()) + if err != nil { + t.Fatalf("Failed to create NATS store: %v", err) + } + defer natsStore.Close() + + handoff := harnessNats.NewHandoffManager(natsStore) + + workflowRunID := fmt.Sprintf("wfr-e2e-%d", time.Now().UnixNano()) + t.Logf("Workflow Run ID: %s", workflowRunID) + + _, err = handoff.StartWorkflow(ctx, harnessNats.StartWorkflowInput{ + WorkflowID: "wf-analyze-fix", + WorkflowRunID: workflowRunID, + GitBranch: "agent/bugfix", + SharedData: map[string]interface{}{ + "target_language": "go", + "fix_strategy": "minimal", + }, + }) + if err != nil { + t.Fatalf("Failed to start workflow: %v", err) + } + + tmpDir, err := os.MkdirTemp("", "multi-agent-e2e-*") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + defer os.RemoveAll(tmpDir) + + buggyCode := `package main + +import "fmt" + +func main() { + numbers := []int{1, 2, 3, 4, 5} + + // BUG: This will cause an index out of bounds panic + for i := 0; i <= len(numbers); i++ { + fmt.Println(numbers[i]) + } + + // BUG: Unused variable + unused := "this variable is never used" + + // BUG: Missing error handling + result := divide(10, 0) + fmt.Println(result) +} + +func divide(a, b int) int { + return a / b +} +` + buggyFilePath := tmpDir + "/buggy.go" + if err := os.WriteFile(buggyFilePath, []byte(buggyCode), 0644); err != nil { + t.Fatalf("Failed to write buggy code: %v", err) + } + + t.Log("") + t.Log("========================================") + t.Log("PHASE 1: ANALYZER AGENT") + t.Log("========================================") + + analyzerRunID := fmt.Sprintf("run-analyzer-%d", time.Now().UnixNano()) + _, _, err = handoff.StartStep(ctx, harnessNats.StartStepInput{ + WorkflowRunID: workflowRunID, + StepName: "analyze", + AgentName: "code-analyzer", + RunID: analyzerRunID, + Task: "Analyze Go code for bugs", + }) + if err != nil { + t.Fatalf("Failed to start analyzer step: %v", err) + } + + ws1 := workspace.NewHostWorkspace(tmpDir) + if err := ws1.Initialize(ctx); err != nil { + t.Fatalf("Failed to initialize workspace for analyzer: %v", err) + } + + toolRegistry1 := tools.NewToolRegistry(genkitApp, tmpDir) + if err := toolRegistry1.RegisterBuiltinTools(); err != nil { + t.Fatalf("Failed to register tools for analyzer: %v", err) + } + + analyzerConfig := harness.DefaultHarnessConfig() + analyzerAgentConfig := &harness.AgentHarnessConfig{ + MaxSteps: 15, + DoomLoopThreshold: 3, + Timeout: 3 * time.Minute, + } + + modelName := formatModelNameForMultiAgent(cfg.AIProvider, cfg.AIModel) + + analyzerPrompt := `You are a code analyzer agent. Your job is to find bugs in code. + +IMPORTANT: You must be thorough but concise. After analyzing, output your findings in a structured format. + +Your task: Read the file buggy.go and identify ALL bugs. For each bug: +1. Line number +2. Bug type (e.g., "index out of bounds", "unused variable", "missing error handling") +3. Brief description + +After reading the file and analyzing it, create a file called "analysis_report.json" with your findings in this exact JSON format: +{ + "bugs": [ + {"line": 9, "type": "index_out_of_bounds", "description": "Loop iterates one past array length"}, + ... + ], + "total_bugs": 3 +} + +Then confirm you've completed the analysis.` + + analyzer := harness.NewAgenticExecutor( + genkitApp, + analyzerConfig, + analyzerAgentConfig, + harness.WithWorkspace(ws1), + harness.WithModelName(modelName), + harness.WithSystemPrompt(analyzerPrompt), + ) + + t.Log("Running Analyzer Agent with real LLM...") + startTime := time.Now() + + analyzerResult, err := analyzer.Execute(ctx, analyzerRunID, "Analyze buggy.go and create analysis_report.json with your findings.", toolRegistry1.All()) + + analyzerDuration := time.Since(startTime) + t.Logf("Analyzer completed in %v", analyzerDuration) + t.Logf("Analyzer result: success=%v, steps=%d, tokens=%d", analyzerResult.Success, analyzerResult.TotalSteps, analyzerResult.TotalTokens) + + if !analyzerResult.Success { + t.Logf("Analyzer error: %s", analyzerResult.Error) + t.Logf("Analyzer response: %s", analyzerResult.Response) + } + + reportPath := tmpDir + "/analysis_report.json" + reportContent, err := os.ReadFile(reportPath) + if err != nil { + t.Logf("Warning: Could not read analysis_report.json: %v", err) + reportContent = []byte(`{"bugs": [{"line": 9, "type": "index_out_of_bounds", "description": "Loop condition should be < not <="}], "total_bugs": 1}`) + } + t.Logf("Analysis report:\n%s", string(reportContent)) + + _, err = natsStore.PutRunFile(ctx, analyzerRunID, "analysis_report.json", bytes.NewReader(reportContent), harnessNats.PutFileOptions{ + ContentType: "application/json", + }) + if err != nil { + t.Fatalf("Failed to store analysis report in NATS: %v", err) + } + t.Log("Analysis report stored in NATS ObjectStore") + + err = handoff.CompleteStep(ctx, harnessNats.CompleteStepInput{ + RunID: analyzerRunID, + WorkflowRunID: workflowRunID, + Status: "completed", + Result: string(reportContent), + Summary: "Code analysis complete - bugs identified", + FilesModified: []string{}, + Commits: []string{}, + }) + if err != nil { + t.Fatalf("Failed to complete analyzer step: %v", err) + } + t.Log("Analyzer step marked complete in NATS KV") + + t.Log("") + t.Log("========================================") + t.Log("PHASE 2: FIXER AGENT") + t.Log("========================================") + + prevCtx, err := handoff.GetPreviousStepContext(ctx, workflowRunID) + if err != nil { + t.Fatalf("Failed to get previous step context: %v", err) + } + t.Logf("Fixer received handoff from: %s (agent: %s)", prevCtx.StepName, prevCtx.AgentName) + t.Logf("Previous step summary: %s", prevCtx.Summary) + + reader, _, err := natsStore.GetRunFile(ctx, analyzerRunID, "analysis_report.json") + if err != nil { + t.Fatalf("Failed to retrieve analysis report from NATS: %v", err) + } + downloadedReport, err := io.ReadAll(reader) + reader.Close() + if err != nil { + t.Fatalf("Failed to read analysis report: %v", err) + } + t.Logf("Fixer downloaded analysis report from NATS: %d bytes", len(downloadedReport)) + + fixerRunID := fmt.Sprintf("run-fixer-%d", time.Now().UnixNano()) + _, _, err = handoff.StartStep(ctx, harnessNats.StartStepInput{ + WorkflowRunID: workflowRunID, + StepName: "fix", + AgentName: "code-fixer", + RunID: fixerRunID, + Task: "Fix bugs identified by analyzer", + }) + if err != nil { + t.Fatalf("Failed to start fixer step: %v", err) + } + + ws2 := workspace.NewHostWorkspace(tmpDir) + if err := ws2.Initialize(ctx); err != nil { + t.Fatalf("Failed to initialize workspace for fixer: %v", err) + } + + // Reuse the same tools from analyzer - they're already registered with GenKit + // GenKit tools are global per-app, so we can't register them twice + + fixerConfig := harness.DefaultHarnessConfig() + fixerAgentConfig := &harness.AgentHarnessConfig{ + MaxSteps: 15, + DoomLoopThreshold: 3, + Timeout: 3 * time.Minute, + } + + fixerPrompt := fmt.Sprintf(`You are a code fixer agent. Your job is to fix bugs identified by the analyzer. + +CONTEXT FROM PREVIOUS AGENT: +The analyzer found these bugs: +%s + +Your task: +1. Read buggy.go +2. Fix the bugs identified in the analysis report +3. Write the fixed code back to buggy.go (overwrite it) +4. Verify the fix by reading the file again + +The fixes should be minimal - only fix what's broken, don't refactor unnecessarily. + +For the index out of bounds bug: change <= to < in the loop condition. +For unused variable: remove it or use it. +For divide by zero: add error handling. + +After fixing, confirm the changes you made.`, string(downloadedReport)) + + fixer := harness.NewAgenticExecutor( + genkitApp, + fixerConfig, + fixerAgentConfig, + harness.WithWorkspace(ws2), + harness.WithModelName(modelName), + harness.WithSystemPrompt(fixerPrompt), + ) + + t.Log("Running Fixer Agent with real LLM...") + startTime = time.Now() + + fixerResult, err := fixer.Execute(ctx, fixerRunID, "Fix the bugs in buggy.go based on the analysis report.", toolRegistry1.All()) + + fixerDuration := time.Since(startTime) + t.Logf("Fixer completed in %v", fixerDuration) + t.Logf("Fixer result: success=%v, steps=%d, tokens=%d", fixerResult.Success, fixerResult.TotalSteps, fixerResult.TotalTokens) + + if !fixerResult.Success { + t.Logf("Fixer error: %s", fixerResult.Error) + t.Logf("Fixer response: %s", fixerResult.Response) + } + + err = handoff.CompleteStep(ctx, harnessNats.CompleteStepInput{ + RunID: fixerRunID, + WorkflowRunID: workflowRunID, + Status: "completed", + Result: "Bugs fixed", + Summary: "Code fixes applied based on analysis", + FilesModified: []string{"buggy.go"}, + Commits: []string{}, + }) + if err != nil { + t.Fatalf("Failed to complete fixer step: %v", err) + } + + t.Log("") + t.Log("========================================") + t.Log("VERIFICATION") + t.Log("========================================") + + fixedCode, err := os.ReadFile(buggyFilePath) + if err != nil { + t.Fatalf("Failed to read fixed code: %v", err) + } + t.Logf("Fixed code:\n%s", string(fixedCode)) + + indexBugFixed := !strings.Contains(string(fixedCode), "i <= len(numbers)") + t.Logf("Index out of bounds bug fixed: %v", indexBugFixed) + + finalCtx, err := natsStore.GetWorkflowContext(ctx, workflowRunID) + if err != nil { + t.Fatalf("Failed to get final workflow context: %v", err) + } + + t.Log("") + t.Log("========================================") + t.Log("FINAL RESULTS") + t.Log("========================================") + t.Logf("Workflow ID: %s", finalCtx.WorkflowID) + t.Logf("Workflow Run ID: %s", finalCtx.WorkflowRunID) + t.Logf("Total Steps: %d", len(finalCtx.Steps)) + for i, step := range finalCtx.Steps { + t.Logf(" Step %d: %s by %s (%s)", i+1, step.StepName, step.AgentName, step.Status) + } + t.Logf("Analyzer Duration: %v", analyzerDuration) + t.Logf("Fixer Duration: %v", fixerDuration) + t.Logf("Total Duration: %v", analyzerDuration+fixerDuration) + t.Logf("Total Tokens: %d (analyzer) + %d (fixer) = %d", analyzerResult.TotalTokens, fixerResult.TotalTokens, analyzerResult.TotalTokens+fixerResult.TotalTokens) + + if len(finalCtx.Steps) != 2 { + t.Errorf("Expected 2 workflow steps, got %d", len(finalCtx.Steps)) + } + + if !analyzerResult.Success { + t.Error("Analyzer agent did not complete successfully") + } + + if !fixerResult.Success { + t.Error("Fixer agent did not complete successfully") + } + + t.Log("") + t.Log("========================================") + t.Log("MULTI-AGENT WORKFLOW COMPLETE!") + t.Log("========================================") + t.Log("Two AI agents successfully collaborated via NATS:") + t.Log(" 1. Analyzer found bugs and stored report in NATS") + t.Log(" 2. Fixer retrieved report from NATS and applied fixes") + t.Log(" 3. Workflow state tracked in NATS KV throughout") +} + +func initGenkitForMultiAgent(ctx context.Context, cfg *config.Config) (*genkit.Genkit, error) { + os.Setenv("OTEL_SDK_DISABLED", "true") + + promptDir := "/tmp/multi-agent-e2e-prompts" + os.MkdirAll(promptDir, 0755) + + provider := strings.ToLower(cfg.AIProvider) + + if strings.HasPrefix(strings.ToLower(cfg.AIModel), "claude") { + provider = "anthropic" + } else if strings.HasPrefix(strings.ToLower(cfg.AIModel), "gemini") { + provider = "gemini" + } else if strings.HasPrefix(strings.ToLower(cfg.AIModel), "gpt") { + provider = "openai" + } + + logging.Info("Initializing GenKit for multi-agent e2e with provider: %s, model: %s", provider, cfg.AIModel) + + switch provider { + case "openai": + return genkit.Init(ctx, + genkit.WithPlugins(&openai.OpenAI{APIKey: cfg.AIAPIKey}), + genkit.WithPromptDir(promptDir)), nil + + case "gemini", "googlegenai": + return genkit.Init(ctx, + genkit.WithPlugins(&googlegenai.GoogleAI{}), + genkit.WithPromptDir(promptDir)), nil + + case "anthropic": + var plugin *anthropic_oauth.AnthropicOAuth + if cfg.AIAuthType == "oauth" && cfg.AIOAuthToken != "" { + plugin = &anthropic_oauth.AnthropicOAuth{OAuthToken: cfg.AIOAuthToken} + } else if cfg.AIAPIKey != "" { + plugin = &anthropic_oauth.AnthropicOAuth{APIKey: cfg.AIAPIKey} + } else { + return nil, fmt.Errorf("Anthropic requires OAuth token or API key") + } + return genkit.Init(ctx, + genkit.WithPlugins(plugin), + genkit.WithPromptDir(promptDir)), nil + + default: + return nil, fmt.Errorf("unsupported provider: %s", provider) + } +} + +func formatModelNameForMultiAgent(provider, model string) string { + provider = strings.ToLower(provider) + + if strings.HasPrefix(strings.ToLower(model), "claude") { + provider = "anthropic" + } else if strings.HasPrefix(strings.ToLower(model), "gemini") { + provider = "googlegenai" + } else if strings.HasPrefix(strings.ToLower(model), "gpt") { + provider = "openai" + } + + return provider + "/" + model +} diff --git a/pkg/harness/nats/handoff.go b/pkg/harness/nats/handoff.go new file mode 100644 index 00000000..d2f17912 --- /dev/null +++ b/pkg/harness/nats/handoff.go @@ -0,0 +1,333 @@ +package nats + +import ( + "bytes" + "context" + "fmt" + "io" + "os" + "path/filepath" + "time" +) + +type HandoffManager struct { + store *Store +} + +func NewHandoffManager(store *Store) *HandoffManager { + return &HandoffManager{store: store} +} + +type StartWorkflowInput struct { + WorkflowID string + WorkflowRunID string + GitBranch string + SharedData map[string]interface{} +} + +func (h *HandoffManager) StartWorkflow(ctx context.Context, input StartWorkflowInput) (*WorkflowContext, error) { + wctx := &WorkflowContext{ + WorkflowID: input.WorkflowID, + WorkflowRunID: input.WorkflowRunID, + StartedAt: time.Now(), + GitBranch: input.GitBranch, + Steps: []WorkflowStepSummary{}, + SharedData: input.SharedData, + } + + if err := h.store.SetWorkflowContext(ctx, wctx); err != nil { + return nil, fmt.Errorf("save workflow context: %w", err) + } + + return wctx, nil +} + +type StartStepInput struct { + WorkflowRunID string + StepName string + AgentName string + RunID string + Task string +} + +func (h *HandoffManager) StartStep(ctx context.Context, input StartStepInput) (*RunState, *WorkflowContext, error) { + wctx, err := h.store.GetWorkflowContext(ctx, input.WorkflowRunID) + if err != nil { + return nil, nil, err + } + if wctx == nil { + return nil, nil, fmt.Errorf("workflow not found: %s", input.WorkflowRunID) + } + + runState := &RunState{ + RunID: input.RunID, + AgentID: input.AgentName, + AgentName: input.AgentName, + WorkflowID: wctx.WorkflowID, + StepName: input.StepName, + Status: "running", + StartedAt: time.Now(), + Task: input.Task, + GitBranch: wctx.GitBranch, + } + + if err := h.store.SetRunState(ctx, runState); err != nil { + return nil, nil, fmt.Errorf("save run state: %w", err) + } + + return runState, wctx, nil +} + +type CompleteStepInput struct { + RunID string + WorkflowRunID string + Status string + Result string + Error string + Summary string + FilesModified []string + Commits []string +} + +func (h *HandoffManager) CompleteStep(ctx context.Context, input CompleteStepInput) error { + runState, err := h.store.GetRunState(ctx, input.RunID) + if err != nil { + return err + } + if runState == nil { + return fmt.Errorf("run not found: %s", input.RunID) + } + + now := time.Now() + runState.Status = input.Status + runState.Result = input.Result + runState.Error = input.Error + runState.CompletedAt = &now + + if err := h.store.SetRunState(ctx, runState); err != nil { + return fmt.Errorf("update run state: %w", err) + } + + stepSummary := WorkflowStepSummary{ + StepName: runState.StepName, + AgentName: runState.AgentName, + RunID: runState.RunID, + Status: input.Status, + StartedAt: runState.StartedAt, + CompletedAt: &now, + Summary: input.Summary, + FilesModified: input.FilesModified, + Commits: input.Commits, + } + + if err := h.store.AddWorkflowStep(ctx, input.WorkflowRunID, stepSummary); err != nil { + return fmt.Errorf("add workflow step: %w", err) + } + + return nil +} + +type PreviousStepContext struct { + StepName string + AgentName string + Status string + Summary string + FilesModified []string + Commits []string + OutputFiles []*FileMetadata +} + +func (h *HandoffManager) GetPreviousStepContext(ctx context.Context, workflowRunID string) (*PreviousStepContext, error) { + wctx, err := h.store.GetWorkflowContext(ctx, workflowRunID) + if err != nil { + return nil, err + } + if wctx == nil || len(wctx.Steps) == 0 { + return nil, nil + } + + lastStep := wctx.Steps[len(wctx.Steps)-1] + + outputFiles, err := h.store.ListRunFiles(ctx, lastStep.RunID) + if err != nil { + outputFiles = nil + } + + return &PreviousStepContext{ + StepName: lastStep.StepName, + AgentName: lastStep.AgentName, + Status: lastStep.Status, + Summary: lastStep.Summary, + FilesModified: lastStep.FilesModified, + Commits: lastStep.Commits, + OutputFiles: outputFiles, + }, nil +} + +func (h *HandoffManager) UploadOutputFile(ctx context.Context, runID, localPath string) (*FileMetadata, error) { + file, err := os.Open(localPath) + if err != nil { + return nil, fmt.Errorf("open file: %w", err) + } + defer file.Close() + + filename := filepath.Base(localPath) + return h.store.PutRunFile(ctx, runID, filename, file, PutFileOptions{ + TTL: 24 * time.Hour, + Metadata: map[string]string{ + "local_path": localPath, + }, + }) +} + +func (h *HandoffManager) DownloadOutputFile(ctx context.Context, runID, filename, localPath string) error { + reader, _, err := h.store.GetRunFile(ctx, runID, filename) + if err != nil { + return err + } + if reader == nil { + return fmt.Errorf("file not found: %s/%s", runID, filename) + } + defer reader.Close() + + if err := os.MkdirAll(filepath.Dir(localPath), 0755); err != nil { + return fmt.Errorf("create directory: %w", err) + } + + file, err := os.Create(localPath) + if err != nil { + return fmt.Errorf("create file: %w", err) + } + defer file.Close() + + if _, err := io.Copy(file, reader); err != nil { + return fmt.Errorf("write file: %w", err) + } + + return nil +} + +func (h *HandoffManager) DownloadPreviousOutputs(ctx context.Context, workflowRunID, localDir string) (int, error) { + prevCtx, err := h.GetPreviousStepContext(ctx, workflowRunID) + if err != nil { + return 0, err + } + if prevCtx == nil || len(prevCtx.OutputFiles) == 0 { + return 0, nil + } + + wctx, err := h.store.GetWorkflowContext(ctx, workflowRunID) + if err != nil { + return 0, err + } + lastStep := wctx.Steps[len(wctx.Steps)-1] + + downloaded := 0 + for _, fileMeta := range prevCtx.OutputFiles { + filename := filepath.Base(fileMeta.Key) + localPath := filepath.Join(localDir, filename) + + if err := h.DownloadOutputFile(ctx, lastStep.RunID, filename, localPath); err != nil { + continue + } + downloaded++ + } + + return downloaded, nil +} + +type SharedDataUpdate struct { + Key string + Value interface{} +} + +func (h *HandoffManager) UpdateSharedData(ctx context.Context, workflowRunID string, updates ...SharedDataUpdate) error { + wctx, err := h.store.GetWorkflowContext(ctx, workflowRunID) + if err != nil { + return err + } + if wctx == nil { + return fmt.Errorf("workflow not found: %s", workflowRunID) + } + + if wctx.SharedData == nil { + wctx.SharedData = make(map[string]interface{}) + } + + for _, update := range updates { + wctx.SharedData[update.Key] = update.Value + } + + return h.store.SetWorkflowContext(ctx, wctx) +} + +func (h *HandoffManager) GetSharedData(ctx context.Context, workflowRunID, key string) (interface{}, error) { + wctx, err := h.store.GetWorkflowContext(ctx, workflowRunID) + if err != nil { + return nil, err + } + if wctx == nil { + return nil, nil + } + + return wctx.SharedData[key], nil +} + +type PreloadFile struct { + Key string + LocalPath string + TTL time.Duration +} + +func (h *HandoffManager) PreloadFiles(ctx context.Context, files []PreloadFile) error { + for _, f := range files { + data, err := os.ReadFile(f.LocalPath) + if err != nil { + return fmt.Errorf("read %s: %w", f.LocalPath, err) + } + + ttl := f.TTL + if ttl == 0 { + ttl = 24 * time.Hour + } + + _, err = h.store.PutSharedFile(ctx, f.Key, bytes.NewReader(data), PutFileOptions{ + TTL: ttl, + Metadata: map[string]string{ + "local_path": f.LocalPath, + }, + }) + if err != nil { + return fmt.Errorf("upload %s: %w", f.Key, err) + } + } + + return nil +} + +func (h *HandoffManager) DownloadPreloadedFile(ctx context.Context, key, localPath string) error { + reader, _, err := h.store.GetSharedFile(ctx, key) + if err != nil { + return err + } + if reader == nil { + return fmt.Errorf("file not found: %s", key) + } + defer reader.Close() + + if err := os.MkdirAll(filepath.Dir(localPath), 0755); err != nil { + return fmt.Errorf("create directory: %w", err) + } + + file, err := os.Create(localPath) + if err != nil { + return fmt.Errorf("create file: %w", err) + } + defer file.Close() + + if _, err := io.Copy(file, reader); err != nil { + return fmt.Errorf("write file: %w", err) + } + + return nil +} diff --git a/pkg/harness/nats/lattice_adapter.go b/pkg/harness/nats/lattice_adapter.go new file mode 100644 index 00000000..5a3b8fed --- /dev/null +++ b/pkg/harness/nats/lattice_adapter.go @@ -0,0 +1,532 @@ +// Package nats provides NATS integration for the agentic harness. +// This adapter bridges the harness with Station's existing lattice infrastructure. +package nats + +import ( + "context" + "fmt" + "io" + "os" + "path/filepath" + "strings" + "sync" + "time" + + "station/internal/lattice/work" + + "github.com/nats-io/nats.go" +) + +// LatticeAdapter bridges the agentic harness with Station's existing lattice infrastructure. +// It uses the existing WorkStore for state management and provides file storage via Object Store. +type LatticeAdapter struct { + workStore *work.WorkStore + fileStore *FileStore + stationID string + mu sync.RWMutex +} + +// LatticeAdapterConfig configures the lattice adapter. +type LatticeAdapterConfig struct { + // StationID identifies this station + StationID string + + // WorkStoreConfig for the lattice work store + WorkStoreConfig work.WorkStoreConfig + + // FileStoreConfig for the object store + FileStoreConfig FileStoreConfig +} + +// FileStoreConfig configures the file object store. +type FileStoreConfig struct { + // Bucket name for the object store + Bucket string + + // MaxFileSize in bytes (default: 100MB) + MaxFileSize int64 + + // MaxBucketSize in bytes (default: 1GB) + MaxBucketSize int64 +} + +// DefaultFileStoreConfig returns sensible defaults. +func DefaultFileStoreConfig() FileStoreConfig { + return FileStoreConfig{ + Bucket: "harness-files", + MaxFileSize: 100 * 1024 * 1024, // 100MB + MaxBucketSize: 1 * 1024 * 1024 * 1024, // 1GB + } +} + +// NewLatticeAdapter creates a new adapter using existing lattice infrastructure. +func NewLatticeAdapter(js nats.JetStreamContext, config LatticeAdapterConfig) (*LatticeAdapter, error) { + // Create or get the existing WorkStore + workStore, err := work.NewWorkStore(js, config.WorkStoreConfig) + if err != nil { + return nil, fmt.Errorf("create work store: %w", err) + } + + // Create file store for harness-specific file sharing + fileStore, err := NewFileStore(js, config.FileStoreConfig) + if err != nil { + return nil, fmt.Errorf("create file store: %w", err) + } + + return &LatticeAdapter{ + workStore: workStore, + fileStore: fileStore, + stationID: config.StationID, + }, nil +} + +// ==================== Work State Operations (using existing WorkStore) ==================== + +func (a *LatticeAdapter) CreateHarnessWork(ctx context.Context, input CreateHarnessWorkInput) (*work.WorkRecord, error) { + contextMap := map[string]string{ + "harness": "agentic", + "workflow_run_id": input.WorkflowRunID, + "step_id": input.StepID, + } + if input.GitBranch != "" { + contextMap["git_branch"] = input.GitBranch + } + + record := &work.WorkRecord{ + WorkID: input.WorkID, + OrchestratorRunID: input.OrchestratorRunID, + ParentWorkID: input.ParentWorkID, + SourceStation: a.stationID, + TargetStation: a.stationID, + AgentID: input.AgentID, + AgentName: input.AgentName, + Task: input.Task, + Context: contextMap, + Status: work.StatusAssigned, + AssignedAt: time.Now(), + } + + if err := a.workStore.Assign(ctx, record); err != nil { + return nil, fmt.Errorf("assign work: %w", err) + } + + return record, nil +} + +type CreateHarnessWorkInput struct { + WorkID string + WorkflowRunID string + StepID string + ParentWorkID string + AgentID string + AgentName string + Task string + OrchestratorRunID string + GitBranch string +} + +// UpdateHarnessWorkStatus updates the status of a harness work record. +func (a *LatticeAdapter) UpdateHarnessWorkStatus(ctx context.Context, workID, status string, result *work.WorkResult) error { + return a.workStore.UpdateStatus(ctx, workID, status, result) +} + +// GetHarnessWork retrieves a harness work record. +func (a *LatticeAdapter) GetHarnessWork(ctx context.Context, workID string) (*work.WorkRecord, error) { + return a.workStore.Get(ctx, workID) +} + +// WatchHarnessWork watches for updates to a specific work record. +func (a *LatticeAdapter) WatchHarnessWork(ctx context.Context, workID string) (<-chan *work.WorkRecord, error) { + return a.workStore.Watch(ctx, workID) +} + +// GetWorkflowWork retrieves all work records for a workflow run. +func (a *LatticeAdapter) GetWorkflowWork(ctx context.Context, orchestratorRunID string) ([]*work.WorkRecord, error) { + return a.workStore.GetByOrchestrator(ctx, orchestratorRunID) +} + +// ==================== File Operations (new functionality) ==================== + +// UploadOutputFile uploads an output file from a harness run. +func (a *LatticeAdapter) UploadOutputFile(ctx context.Context, workID, localPath string) (*FileMetadata, error) { + file, err := os.Open(localPath) + if err != nil { + return nil, fmt.Errorf("open file: %w", err) + } + defer file.Close() + + filename := filepath.Base(localPath) + key := fmt.Sprintf("work/%s/output/%s", workID, filename) + + return a.fileStore.Put(ctx, key, file, PutFileOptions{ + TTL: 24 * time.Hour, + Metadata: map[string]string{ + "work_id": workID, + "local_path": localPath, + }, + }) +} + +// DownloadOutputFile downloads an output file from a harness run. +func (a *LatticeAdapter) DownloadOutputFile(ctx context.Context, workID, filename, localPath string) error { + key := fmt.Sprintf("work/%s/output/%s", workID, filename) + + reader, _, err := a.fileStore.Get(ctx, key) + if err != nil { + return err + } + if reader == nil { + return fmt.Errorf("file not found: %s", key) + } + defer reader.Close() + + if err := os.MkdirAll(filepath.Dir(localPath), 0755); err != nil { + return fmt.Errorf("create directory: %w", err) + } + + file, err := os.Create(localPath) + if err != nil { + return fmt.Errorf("create file: %w", err) + } + defer file.Close() + + if _, err := io.Copy(file, reader); err != nil { + return fmt.Errorf("write file: %w", err) + } + + return nil +} + +// ListOutputFiles lists output files from a harness run. +func (a *LatticeAdapter) ListOutputFiles(ctx context.Context, workID string) ([]*FileMetadata, error) { + prefix := fmt.Sprintf("work/%s/output/", workID) + return a.fileStore.List(ctx, prefix) +} + +// UploadSharedFile uploads a file shared across workflow steps. +func (a *LatticeAdapter) UploadSharedFile(ctx context.Context, workflowRunID, key string, reader io.Reader, opts PutFileOptions) (*FileMetadata, error) { + fullKey := fmt.Sprintf("workflow/%s/shared/%s", workflowRunID, key) + return a.fileStore.Put(ctx, fullKey, reader, opts) +} + +// DownloadSharedFile downloads a shared file. +func (a *LatticeAdapter) DownloadSharedFile(ctx context.Context, workflowRunID, key, localPath string) error { + fullKey := fmt.Sprintf("workflow/%s/shared/%s", workflowRunID, key) + + reader, _, err := a.fileStore.Get(ctx, fullKey) + if err != nil { + return err + } + if reader == nil { + return fmt.Errorf("file not found: %s", fullKey) + } + defer reader.Close() + + if err := os.MkdirAll(filepath.Dir(localPath), 0755); err != nil { + return fmt.Errorf("create directory: %w", err) + } + + file, err := os.Create(localPath) + if err != nil { + return fmt.Errorf("create file: %w", err) + } + defer file.Close() + + if _, err := io.Copy(file, reader); err != nil { + return fmt.Errorf("write file: %w", err) + } + + return nil +} + +// ==================== Handoff Operations ==================== + +// GetPreviousStepFiles gets output files from the previous workflow step. +func (a *LatticeAdapter) GetPreviousStepFiles(ctx context.Context, orchestratorRunID string) ([]*FileMetadata, error) { + // Get all work records for the workflow + records, err := a.workStore.GetByOrchestrator(ctx, orchestratorRunID) + if err != nil { + return nil, err + } + + if len(records) == 0 { + return nil, nil + } + + // Find the most recently completed step + var lastCompleted *work.WorkRecord + for _, r := range records { + if r.Status == work.StatusComplete { + if lastCompleted == nil || r.CompletedAt.After(lastCompleted.CompletedAt) { + lastCompleted = r + } + } + } + + if lastCompleted == nil { + return nil, nil + } + + // List output files from that step + return a.ListOutputFiles(ctx, lastCompleted.WorkID) +} + +// DownloadPreviousStepFiles downloads all output files from the previous step to a local directory. +func (a *LatticeAdapter) DownloadPreviousStepFiles(ctx context.Context, orchestratorRunID, localDir string) (int, error) { + files, err := a.GetPreviousStepFiles(ctx, orchestratorRunID) + if err != nil { + return 0, err + } + + downloaded := 0 + for _, fileMeta := range files { + // Extract filename from key (work/{id}/output/{filename}) + parts := strings.Split(fileMeta.Key, "/") + if len(parts) < 4 { + continue + } + filename := parts[len(parts)-1] + workID := parts[1] + + localPath := filepath.Join(localDir, filename) + if err := a.DownloadOutputFile(ctx, workID, filename, localPath); err != nil { + continue + } + downloaded++ + } + + return downloaded, nil +} + +// Close releases resources. +func (a *LatticeAdapter) Close() error { + return a.fileStore.Close() +} + +// ==================== FileStore (Object Store wrapper) ==================== + +// FileStore provides object storage for harness files. +type FileStore struct { + js nats.JetStreamContext + objectStore nats.ObjectStore + config FileStoreConfig + mu sync.RWMutex +} + +// NewFileStore creates a new file store. +func NewFileStore(js nats.JetStreamContext, config FileStoreConfig) (*FileStore, error) { + if config.Bucket == "" { + config.Bucket = "harness-files" + } + if config.MaxFileSize == 0 { + config.MaxFileSize = 100 * 1024 * 1024 + } + if config.MaxBucketSize == 0 { + config.MaxBucketSize = 1 * 1024 * 1024 * 1024 + } + + return &FileStore{ + js: js, + config: config, + }, nil +} + +func (s *FileStore) getObjectStore(ctx context.Context) (nats.ObjectStore, error) { + s.mu.RLock() + if s.objectStore != nil { + os := s.objectStore + s.mu.RUnlock() + return os, nil + } + s.mu.RUnlock() + + s.mu.Lock() + defer s.mu.Unlock() + + if s.objectStore != nil { + return s.objectStore, nil + } + + os, err := s.js.ObjectStore(s.config.Bucket) + if err == nil { + s.objectStore = os + return os, nil + } + + if err == nats.ErrBucketNotFound || err == nats.ErrStreamNotFound || strings.Contains(err.Error(), "stream not found") { + os, err = s.js.CreateObjectStore(&nats.ObjectStoreConfig{ + Bucket: s.config.Bucket, + Description: "Agentic harness file storage", + MaxBytes: s.config.MaxBucketSize, + }) + if err != nil { + return nil, fmt.Errorf("create object store %s: %w", s.config.Bucket, err) + } + s.objectStore = os + return os, nil + } + + return nil, fmt.Errorf("get object store %s: %w", s.config.Bucket, err) +} + +// Put stores a file. +func (s *FileStore) Put(ctx context.Context, key string, reader io.Reader, opts PutFileOptions) (*FileMetadata, error) { + os, err := s.getObjectStore(ctx) + if err != nil { + return nil, err + } + + data, err := io.ReadAll(reader) + if err != nil { + return nil, fmt.Errorf("read file data: %w", err) + } + + if s.config.MaxFileSize > 0 && int64(len(data)) > s.config.MaxFileSize { + return nil, fmt.Errorf("file too large: %d bytes (max: %d)", len(data), s.config.MaxFileSize) + } + + meta := &nats.ObjectMeta{ + Name: key, + Description: opts.Description, + Headers: make(nats.Header), + } + + if opts.ContentType != "" { + meta.Headers.Set("Content-Type", opts.ContentType) + } + + if opts.TTL > 0 { + expiresAt := time.Now().Add(opts.TTL) + meta.Headers.Set("X-Expires-At", expiresAt.Format(time.RFC3339)) + } + + for k, v := range opts.Metadata { + meta.Headers.Set("X-Meta-"+k, v) + } + + _, err = os.Put(meta, strings.NewReader(string(data))) + if err != nil { + return nil, fmt.Errorf("put file %s: %w", key, err) + } + + info := &FileMetadata{ + Key: key, + Size: int64(len(data)), + ContentType: opts.ContentType, + CreatedAt: time.Now(), + Metadata: opts.Metadata, + } + + if opts.TTL > 0 { + expires := time.Now().Add(opts.TTL) + info.ExpiresAt = &expires + } + + return info, nil +} + +// Get retrieves a file. +func (s *FileStore) Get(ctx context.Context, key string) (io.ReadCloser, *FileMetadata, error) { + os, err := s.getObjectStore(ctx) + if err != nil { + return nil, nil, err + } + + result, err := os.Get(key) + if err != nil { + if err == nats.ErrObjectNotFound { + return nil, nil, nil + } + return nil, nil, fmt.Errorf("get file %s: %w", key, err) + } + + objInfo, err := result.Info() + if err != nil { + result.Close() + return nil, nil, fmt.Errorf("get file info %s: %w", key, err) + } + + info := &FileMetadata{ + Key: objInfo.Name, + Size: int64(objInfo.Size), + CreatedAt: objInfo.ModTime, + Metadata: make(map[string]string), + } + + if objInfo.Headers != nil { + info.ContentType = objInfo.Headers.Get("Content-Type") + + if expiresStr := objInfo.Headers.Get("X-Expires-At"); expiresStr != "" { + if t, err := time.Parse(time.RFC3339, expiresStr); err == nil { + info.ExpiresAt = &t + } + } + + for k, values := range objInfo.Headers { + if strings.HasPrefix(k, "X-Meta-") && len(values) > 0 { + metaKey := strings.TrimPrefix(k, "X-Meta-") + info.Metadata[metaKey] = values[0] + } + } + } + + return result, info, nil +} + +// Delete removes a file. +func (s *FileStore) Delete(ctx context.Context, key string) error { + os, err := s.getObjectStore(ctx) + if err != nil { + return err + } + + if err := os.Delete(key); err != nil && err != nats.ErrObjectNotFound { + return fmt.Errorf("delete file %s: %w", key, err) + } + + return nil +} + +// List lists files with the given prefix. +func (s *FileStore) List(ctx context.Context, prefix string) ([]*FileMetadata, error) { + os, err := s.getObjectStore(ctx) + if err != nil { + return nil, err + } + + objects, err := os.List() + if err != nil { + return nil, fmt.Errorf("list files: %w", err) + } + + var files []*FileMetadata + for _, obj := range objects { + if prefix == "" || strings.HasPrefix(obj.Name, prefix) { + info := &FileMetadata{ + Key: obj.Name, + Size: int64(obj.Size), + CreatedAt: obj.ModTime, + Metadata: make(map[string]string), + } + + if obj.Headers != nil { + info.ContentType = obj.Headers.Get("Content-Type") + for k, values := range obj.Headers { + if strings.HasPrefix(k, "X-Meta-") && len(values) > 0 { + metaKey := strings.TrimPrefix(k, "X-Meta-") + info.Metadata[metaKey] = values[0] + } + } + } + + files = append(files, info) + } + } + + return files, nil +} + +// Close releases resources. +func (s *FileStore) Close() error { + return nil +} diff --git a/pkg/harness/nats/lattice_integration_test.go b/pkg/harness/nats/lattice_integration_test.go new file mode 100644 index 00000000..1ee52953 --- /dev/null +++ b/pkg/harness/nats/lattice_integration_test.go @@ -0,0 +1,1086 @@ +//go:build integration + +// Package nats provides NATS JetStream integration tests using the lattice embedded server. +// These tests verify that the harness NATS store works correctly with real NATS infrastructure. +// +// Run with: go test ./pkg/harness/nats/... -tags=integration -v +package nats + +import ( + "bytes" + "context" + "fmt" + "io" + "net" + "sync" + "testing" + "time" + + "station/internal/config" + "station/internal/lattice" +) + +// getFreePort finds an available TCP port for testing. +func getFreePort(t *testing.T) int { + t.Helper() + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("Failed to find free port: %v", err) + } + port := listener.Addr().(*net.TCPAddr).Port + listener.Close() + return port +} + +// testEnv holds the test environment with embedded NATS. +type testEnv struct { + server *lattice.EmbeddedServer + client *lattice.Client + store *Store +} + +// setupTestEnv creates a test environment with embedded NATS server. +func setupTestEnv(t *testing.T) *testEnv { + t.Helper() + + port := getFreePort(t) + httpPort := getFreePort(t) + + serverCfg := config.LatticeEmbeddedNATSConfig{ + Port: port, + HTTPPort: httpPort, + StoreDir: t.TempDir(), + } + + server := lattice.NewEmbeddedServer(serverCfg) + if err := server.Start(); err != nil { + t.Fatalf("Failed to start embedded NATS server: %v", err) + } + + clientCfg := config.LatticeConfig{ + StationID: "test-harness", + StationName: "Test Harness Station", + NATS: config.LatticeNATSConfig{URL: server.ClientURL()}, + } + + client, err := lattice.NewClient(clientCfg) + if err != nil { + server.Shutdown() + t.Fatalf("Failed to create lattice client: %v", err) + } + + if err := client.Connect(); err != nil { + server.Shutdown() + t.Fatalf("Failed to connect to NATS: %v", err) + } + + // Create harness store using the lattice client's NATS connection + store, err := NewStore(client.Conn(), DefaultStoreConfig()) + if err != nil { + client.Close() + server.Shutdown() + t.Fatalf("Failed to create harness store: %v", err) + } + + return &testEnv{ + server: server, + client: client, + store: store, + } +} + +// cleanup tears down the test environment. +func (e *testEnv) cleanup() { + if e.store != nil { + e.store.Close() + } + if e.client != nil { + e.client.Close() + } + if e.server != nil { + e.server.Shutdown() + } +} + +// ============================================================================= +// KV Store Tests +// ============================================================================= + +func TestIntegration_KV_BasicOperations(t *testing.T) { + env := setupTestEnv(t) + defer env.cleanup() + + ctx := context.Background() + + t.Run("SetAndGetState", func(t *testing.T) { + key := "test.key.1" + value := []byte("hello world") + + // Set state + err := env.store.SetState(ctx, key, value) + if err != nil { + t.Fatalf("SetState failed: %v", err) + } + + // Get state + got, err := env.store.GetState(ctx, key) + if err != nil { + t.Fatalf("GetState failed: %v", err) + } + + if !bytes.Equal(got, value) { + t.Errorf("GetState returned wrong value: got %q, want %q", got, value) + } + }) + + t.Run("GetNonExistentKey", func(t *testing.T) { + got, err := env.store.GetState(ctx, "nonexistent.key") + if err != nil { + t.Fatalf("GetState failed: %v", err) + } + + if got != nil { + t.Errorf("Expected nil for non-existent key, got %q", got) + } + }) + + t.Run("DeleteState", func(t *testing.T) { + key := "test.delete.key" + value := []byte("to be deleted") + + // Set state + err := env.store.SetState(ctx, key, value) + if err != nil { + t.Fatalf("SetState failed: %v", err) + } + + // Delete state + err = env.store.DeleteState(ctx, key) + if err != nil { + t.Fatalf("DeleteState failed: %v", err) + } + + // Verify deletion + got, err := env.store.GetState(ctx, key) + if err != nil { + t.Fatalf("GetState after delete failed: %v", err) + } + + if got != nil { + t.Errorf("Expected nil after delete, got %q", got) + } + }) + + t.Run("SetAndGetJSON", func(t *testing.T) { + key := "test.json.key" + original := RunState{ + RunID: "run-123", + AgentID: "agent-1", + AgentName: "test-agent", + Status: "running", + Task: "Do something", + StartedAt: time.Now().Truncate(time.Second), + } + + // Set JSON + err := env.store.SetJSON(ctx, key, original) + if err != nil { + t.Fatalf("SetJSON failed: %v", err) + } + + // Get JSON + var got RunState + err = env.store.GetJSON(ctx, key, &got) + if err != nil { + t.Fatalf("GetJSON failed: %v", err) + } + + if got.RunID != original.RunID { + t.Errorf("RunID mismatch: got %q, want %q", got.RunID, original.RunID) + } + if got.AgentName != original.AgentName { + t.Errorf("AgentName mismatch: got %q, want %q", got.AgentName, original.AgentName) + } + if got.Status != original.Status { + t.Errorf("Status mismatch: got %q, want %q", got.Status, original.Status) + } + }) +} + +func TestIntegration_KV_RunState(t *testing.T) { + env := setupTestEnv(t) + defer env.cleanup() + + ctx := context.Background() + + t.Run("SetAndGetRunState", func(t *testing.T) { + state := &RunState{ + RunID: "run-456", + AgentID: "agent-2", + AgentName: "coder", + WorkflowID: "wf-1", + StepName: "analyze", + Status: "running", + StartedAt: time.Now(), + Task: "Analyze the codebase", + GitBranch: "feature/test", + } + + err := env.store.SetRunState(ctx, state) + if err != nil { + t.Fatalf("SetRunState failed: %v", err) + } + + got, err := env.store.GetRunState(ctx, state.RunID) + if err != nil { + t.Fatalf("GetRunState failed: %v", err) + } + + if got == nil { + t.Fatal("GetRunState returned nil") + } + + if got.RunID != state.RunID { + t.Errorf("RunID mismatch: got %q, want %q", got.RunID, state.RunID) + } + if got.WorkflowID != state.WorkflowID { + t.Errorf("WorkflowID mismatch: got %q, want %q", got.WorkflowID, state.WorkflowID) + } + if got.GitBranch != state.GitBranch { + t.Errorf("GitBranch mismatch: got %q, want %q", got.GitBranch, state.GitBranch) + } + }) + + t.Run("UpdateRunState", func(t *testing.T) { + state := &RunState{ + RunID: "run-789", + AgentName: "reviewer", + Status: "running", + StartedAt: time.Now(), + Task: "Review code", + } + + // Create initial state + err := env.store.SetRunState(ctx, state) + if err != nil { + t.Fatalf("SetRunState failed: %v", err) + } + + // Update state + now := time.Now() + state.Status = "completed" + state.Result = "Code looks good!" + state.CompletedAt = &now + + err = env.store.SetRunState(ctx, state) + if err != nil { + t.Fatalf("SetRunState (update) failed: %v", err) + } + + // Verify update + got, err := env.store.GetRunState(ctx, state.RunID) + if err != nil { + t.Fatalf("GetRunState failed: %v", err) + } + + if got.Status != "completed" { + t.Errorf("Status not updated: got %q, want %q", got.Status, "completed") + } + if got.Result != "Code looks good!" { + t.Errorf("Result not updated: got %q", got.Result) + } + if got.CompletedAt == nil { + t.Error("CompletedAt not set") + } + }) +} + +func TestIntegration_KV_WorkflowContext(t *testing.T) { + env := setupTestEnv(t) + defer env.cleanup() + + ctx := context.Background() + + t.Run("SetAndGetWorkflowContext", func(t *testing.T) { + wctx := &WorkflowContext{ + WorkflowID: "wf-analyze-fix", + WorkflowRunID: "wfr-001", + StartedAt: time.Now(), + GitBranch: "agent/wf-analyze-fix", + Steps: []WorkflowStepSummary{}, + SharedData: map[string]interface{}{ + "target_file": "main.go", + "priority": "high", + }, + } + + err := env.store.SetWorkflowContext(ctx, wctx) + if err != nil { + t.Fatalf("SetWorkflowContext failed: %v", err) + } + + got, err := env.store.GetWorkflowContext(ctx, wctx.WorkflowRunID) + if err != nil { + t.Fatalf("GetWorkflowContext failed: %v", err) + } + + if got == nil { + t.Fatal("GetWorkflowContext returned nil") + } + + if got.WorkflowID != wctx.WorkflowID { + t.Errorf("WorkflowID mismatch: got %q, want %q", got.WorkflowID, wctx.WorkflowID) + } + if got.GitBranch != wctx.GitBranch { + t.Errorf("GitBranch mismatch: got %q, want %q", got.GitBranch, wctx.GitBranch) + } + if got.SharedData["target_file"] != "main.go" { + t.Errorf("SharedData mismatch: got %v", got.SharedData) + } + }) + + t.Run("AddWorkflowStep", func(t *testing.T) { + wctx := &WorkflowContext{ + WorkflowID: "wf-multi-step", + WorkflowRunID: "wfr-002", + StartedAt: time.Now(), + GitBranch: "agent/multi-step", + Steps: []WorkflowStepSummary{}, + } + + err := env.store.SetWorkflowContext(ctx, wctx) + if err != nil { + t.Fatalf("SetWorkflowContext failed: %v", err) + } + + // Add first step + now := time.Now() + step1 := WorkflowStepSummary{ + StepName: "analyze", + AgentName: "analyzer", + RunID: "run-step1", + Status: "completed", + StartedAt: now, + CompletedAt: &now, + Summary: "Found 3 issues", + FilesModified: []string{"main.go", "utils.go"}, + Commits: []string{"abc123"}, + } + + err = env.store.AddWorkflowStep(ctx, wctx.WorkflowRunID, step1) + if err != nil { + t.Fatalf("AddWorkflowStep failed: %v", err) + } + + // Add second step + step2 := WorkflowStepSummary{ + StepName: "fix", + AgentName: "fixer", + RunID: "run-step2", + Status: "completed", + StartedAt: now, + Summary: "Fixed all issues", + } + + err = env.store.AddWorkflowStep(ctx, wctx.WorkflowRunID, step2) + if err != nil { + t.Fatalf("AddWorkflowStep (second) failed: %v", err) + } + + // Verify steps + got, err := env.store.GetWorkflowContext(ctx, wctx.WorkflowRunID) + if err != nil { + t.Fatalf("GetWorkflowContext failed: %v", err) + } + + if len(got.Steps) != 2 { + t.Fatalf("Expected 2 steps, got %d", len(got.Steps)) + } + + if got.Steps[0].StepName != "analyze" { + t.Errorf("First step name mismatch: got %q", got.Steps[0].StepName) + } + if got.Steps[1].StepName != "fix" { + t.Errorf("Second step name mismatch: got %q", got.Steps[1].StepName) + } + }) +} + +// ============================================================================= +// ObjectStore Tests +// ============================================================================= + +func TestIntegration_ObjectStore_BasicOperations(t *testing.T) { + env := setupTestEnv(t) + defer env.cleanup() + + ctx := context.Background() + + t.Run("PutAndGetFile", func(t *testing.T) { + content := []byte("Hello from the harness ObjectStore!") + key := "test/hello.txt" + + meta, err := env.store.PutFile(ctx, key, bytes.NewReader(content), PutFileOptions{ + ContentType: "text/plain", + Description: "Test file", + }) + if err != nil { + t.Fatalf("PutFile failed: %v", err) + } + + if meta.Key != key { + t.Errorf("Key mismatch: got %q, want %q", meta.Key, key) + } + if meta.Size != int64(len(content)) { + t.Errorf("Size mismatch: got %d, want %d", meta.Size, len(content)) + } + + // Get file + reader, gotMeta, err := env.store.GetFile(ctx, key) + if err != nil { + t.Fatalf("GetFile failed: %v", err) + } + defer reader.Close() + + gotContent, err := io.ReadAll(reader) + if err != nil { + t.Fatalf("ReadAll failed: %v", err) + } + + if !bytes.Equal(gotContent, content) { + t.Errorf("Content mismatch: got %q, want %q", gotContent, content) + } + + if gotMeta.Size != int64(len(content)) { + t.Errorf("Metadata size mismatch: got %d, want %d", gotMeta.Size, len(content)) + } + }) + + t.Run("GetNonExistentFile", func(t *testing.T) { + reader, meta, err := env.store.GetFile(ctx, "nonexistent/file.txt") + if err != nil { + t.Fatalf("GetFile failed: %v", err) + } + + if reader != nil || meta != nil { + t.Error("Expected nil for non-existent file") + if reader != nil { + reader.Close() + } + } + }) + + t.Run("DeleteFile", func(t *testing.T) { + content := []byte("Delete me") + key := "test/delete-me.txt" + + _, err := env.store.PutFile(ctx, key, bytes.NewReader(content), PutFileOptions{}) + if err != nil { + t.Fatalf("PutFile failed: %v", err) + } + + err = env.store.DeleteFile(ctx, key) + if err != nil { + t.Fatalf("DeleteFile failed: %v", err) + } + + reader, _, err := env.store.GetFile(ctx, key) + if err != nil { + t.Fatalf("GetFile after delete failed: %v", err) + } + + if reader != nil { + reader.Close() + t.Error("Expected nil after delete") + } + }) + + t.Run("ListFiles", func(t *testing.T) { + // Create several files with a common prefix + prefix := "listtest/" + files := []string{"a.txt", "b.txt", "c.txt"} + + for _, f := range files { + key := prefix + f + _, err := env.store.PutFile(ctx, key, bytes.NewReader([]byte(f)), PutFileOptions{}) + if err != nil { + t.Fatalf("PutFile %s failed: %v", key, err) + } + } + + // List files + list, err := env.store.ListFiles(ctx, prefix) + if err != nil { + t.Fatalf("ListFiles failed: %v", err) + } + + if len(list) < len(files) { + t.Errorf("Expected at least %d files, got %d", len(files), len(list)) + } + }) +} + +func TestIntegration_ObjectStore_RunFiles(t *testing.T) { + env := setupTestEnv(t) + defer env.cleanup() + + ctx := context.Background() + runID := "run-output-test" + + t.Run("PutAndGetRunFile", func(t *testing.T) { + content := []byte(`{"analysis": "done", "issues": 3}`) + filename := "analysis.json" + + meta, err := env.store.PutRunFile(ctx, runID, filename, bytes.NewReader(content), PutFileOptions{ + ContentType: "application/json", + Metadata: map[string]string{ + "step": "analyze", + }, + }) + if err != nil { + t.Fatalf("PutRunFile failed: %v", err) + } + + expectedKey := fmt.Sprintf("run/%s/output/%s", runID, filename) + if meta.Key != expectedKey { + t.Errorf("Key mismatch: got %q, want %q", meta.Key, expectedKey) + } + + // Get run file + reader, _, err := env.store.GetRunFile(ctx, runID, filename) + if err != nil { + t.Fatalf("GetRunFile failed: %v", err) + } + defer reader.Close() + + gotContent, err := io.ReadAll(reader) + if err != nil { + t.Fatalf("ReadAll failed: %v", err) + } + + if !bytes.Equal(gotContent, content) { + t.Errorf("Content mismatch: got %q, want %q", gotContent, content) + } + }) + + t.Run("ListRunFiles", func(t *testing.T) { + // Add more files to the run + files := []string{"output1.txt", "output2.txt", "results.json"} + for _, f := range files { + _, err := env.store.PutRunFile(ctx, runID, f, bytes.NewReader([]byte(f)), PutFileOptions{}) + if err != nil { + t.Fatalf("PutRunFile %s failed: %v", f, err) + } + } + + list, err := env.store.ListRunFiles(ctx, runID) + if err != nil { + t.Fatalf("ListRunFiles failed: %v", err) + } + + // Should have at least the files we created (plus analysis.json from previous test) + if len(list) < len(files) { + t.Errorf("Expected at least %d files, got %d", len(files), len(list)) + } + }) +} + +func TestIntegration_ObjectStore_SharedFiles(t *testing.T) { + env := setupTestEnv(t) + defer env.cleanup() + + ctx := context.Background() + + t.Run("PutAndGetSharedFile", func(t *testing.T) { + content := []byte("Shared configuration data") + key := "config/shared.yaml" + + meta, err := env.store.PutSharedFile(ctx, key, bytes.NewReader(content), PutFileOptions{ + ContentType: "application/yaml", + TTL: 24 * time.Hour, + }) + if err != nil { + t.Fatalf("PutSharedFile failed: %v", err) + } + + expectedKey := "shared/" + key + if meta.Key != expectedKey { + t.Errorf("Key mismatch: got %q, want %q", meta.Key, expectedKey) + } + + // Get shared file + reader, _, err := env.store.GetSharedFile(ctx, key) + if err != nil { + t.Fatalf("GetSharedFile failed: %v", err) + } + defer reader.Close() + + gotContent, err := io.ReadAll(reader) + if err != nil { + t.Fatalf("ReadAll failed: %v", err) + } + + if !bytes.Equal(gotContent, content) { + t.Errorf("Content mismatch: got %q, want %q", gotContent, content) + } + }) +} + +// ============================================================================= +// Workflow Handoff Tests +// ============================================================================= + +func TestIntegration_WorkflowHandoff(t *testing.T) { + env := setupTestEnv(t) + defer env.cleanup() + + ctx := context.Background() + handoff := NewHandoffManager(env.store) + + t.Run("FullWorkflowHandoff", func(t *testing.T) { + // Step 1: Start workflow + wctx, err := handoff.StartWorkflow(ctx, StartWorkflowInput{ + WorkflowID: "wf-code-review", + WorkflowRunID: "wfr-handoff-test", + GitBranch: "agent/code-review", + SharedData: map[string]interface{}{ + "target_pr": "PR-123", + "repo": "myorg/myrepo", + }, + }) + if err != nil { + t.Fatalf("StartWorkflow failed: %v", err) + } + + if wctx.WorkflowID != "wf-code-review" { + t.Errorf("WorkflowID mismatch: got %q", wctx.WorkflowID) + } + + // Step 2: Start first step (analyzer) + runState, gotWctx, err := handoff.StartStep(ctx, StartStepInput{ + WorkflowRunID: wctx.WorkflowRunID, + StepName: "analyze", + AgentName: "code-analyzer", + RunID: "run-analyze-1", + Task: "Analyze PR-123 for issues", + }) + if err != nil { + t.Fatalf("StartStep failed: %v", err) + } + + if runState.Status != "running" { + t.Errorf("RunState status mismatch: got %q", runState.Status) + } + if gotWctx.SharedData["target_pr"] != "PR-123" { + t.Errorf("SharedData not preserved: got %v", gotWctx.SharedData) + } + + // Step 3: Complete first step + err = handoff.CompleteStep(ctx, CompleteStepInput{ + RunID: runState.RunID, + WorkflowRunID: wctx.WorkflowRunID, + Status: "completed", + Result: "Found 2 issues: unused variable, missing error check", + Summary: "Code analysis complete", + FilesModified: []string{"main.go"}, + Commits: []string{"def456"}, + }) + if err != nil { + t.Fatalf("CompleteStep failed: %v", err) + } + + // Step 4: Get previous step context (for next agent) + prevCtx, err := handoff.GetPreviousStepContext(ctx, wctx.WorkflowRunID) + if err != nil { + t.Fatalf("GetPreviousStepContext failed: %v", err) + } + + if prevCtx == nil { + t.Fatal("GetPreviousStepContext returned nil") + } + + if prevCtx.StepName != "analyze" { + t.Errorf("Previous step name mismatch: got %q", prevCtx.StepName) + } + if prevCtx.Summary != "Code analysis complete" { + t.Errorf("Previous step summary mismatch: got %q", prevCtx.Summary) + } + + // Step 5: Update shared data + err = handoff.UpdateSharedData(ctx, wctx.WorkflowRunID, + SharedDataUpdate{Key: "issues_found", Value: 2}, + SharedDataUpdate{Key: "analyzer_result", Value: "needs_fix"}, + ) + if err != nil { + t.Fatalf("UpdateSharedData failed: %v", err) + } + + // Step 6: Get shared data + issuesFound, err := handoff.GetSharedData(ctx, wctx.WorkflowRunID, "issues_found") + if err != nil { + t.Fatalf("GetSharedData failed: %v", err) + } + + // Note: JSON unmarshaling converts numbers to float64 + if issuesFound != float64(2) { + t.Errorf("SharedData issues_found mismatch: got %v (type: %T)", issuesFound, issuesFound) + } + + // Step 7: Start second step (fixer) + runState2, _, err := handoff.StartStep(ctx, StartStepInput{ + WorkflowRunID: wctx.WorkflowRunID, + StepName: "fix", + AgentName: "code-fixer", + RunID: "run-fix-1", + Task: "Fix the 2 issues found by analyzer", + }) + if err != nil { + t.Fatalf("StartStep (fixer) failed: %v", err) + } + + // Step 8: Complete second step + err = handoff.CompleteStep(ctx, CompleteStepInput{ + RunID: runState2.RunID, + WorkflowRunID: wctx.WorkflowRunID, + Status: "completed", + Result: "Fixed all issues", + Summary: "Code fixes applied", + FilesModified: []string{"main.go", "utils.go"}, + Commits: []string{"ghi789"}, + }) + if err != nil { + t.Fatalf("CompleteStep (fixer) failed: %v", err) + } + + // Verify final workflow state + finalWctx, err := env.store.GetWorkflowContext(ctx, wctx.WorkflowRunID) + if err != nil { + t.Fatalf("GetWorkflowContext failed: %v", err) + } + + if len(finalWctx.Steps) != 2 { + t.Fatalf("Expected 2 steps, got %d", len(finalWctx.Steps)) + } + + if finalWctx.Steps[0].StepName != "analyze" { + t.Errorf("First step name mismatch: got %q", finalWctx.Steps[0].StepName) + } + if finalWctx.Steps[1].StepName != "fix" { + t.Errorf("Second step name mismatch: got %q", finalWctx.Steps[1].StepName) + } + + t.Logf("Workflow completed successfully with %d steps", len(finalWctx.Steps)) + }) +} + +// ============================================================================= +// Concurrent Access Tests +// ============================================================================= + +func TestIntegration_ConcurrentAccess(t *testing.T) { + env := setupTestEnv(t) + defer env.cleanup() + + ctx := context.Background() + + t.Run("ConcurrentKVWrites", func(t *testing.T) { + numGoroutines := 10 + numWritesPerGoroutine := 20 + + var wg sync.WaitGroup + errors := make(chan error, numGoroutines*numWritesPerGoroutine) + + for i := 0; i < numGoroutines; i++ { + wg.Add(1) + go func(goroutineID int) { + defer wg.Done() + for j := 0; j < numWritesPerGoroutine; j++ { + key := fmt.Sprintf("concurrent.%d.%d", goroutineID, j) + value := []byte(fmt.Sprintf("value-%d-%d", goroutineID, j)) + + if err := env.store.SetState(ctx, key, value); err != nil { + errors <- fmt.Errorf("goroutine %d, write %d: %w", goroutineID, j, err) + return + } + + got, err := env.store.GetState(ctx, key) + if err != nil { + errors <- fmt.Errorf("goroutine %d, read %d: %w", goroutineID, j, err) + return + } + + if !bytes.Equal(got, value) { + errors <- fmt.Errorf("goroutine %d, mismatch %d: got %q, want %q", goroutineID, j, got, value) + return + } + } + }(i) + } + + wg.Wait() + close(errors) + + var errList []error + for err := range errors { + errList = append(errList, err) + } + + if len(errList) > 0 { + t.Fatalf("Concurrent access errors: %v", errList) + } + + t.Logf("Successfully completed %d concurrent writes", numGoroutines*numWritesPerGoroutine) + }) + + t.Run("ConcurrentFileUploads", func(t *testing.T) { + numGoroutines := 5 + numFilesPerGoroutine := 10 + + var wg sync.WaitGroup + errors := make(chan error, numGoroutines*numFilesPerGoroutine) + + for i := 0; i < numGoroutines; i++ { + wg.Add(1) + go func(goroutineID int) { + defer wg.Done() + for j := 0; j < numFilesPerGoroutine; j++ { + key := fmt.Sprintf("concurrent-files/%d/file-%d.txt", goroutineID, j) + content := []byte(fmt.Sprintf("Content from goroutine %d, file %d", goroutineID, j)) + + _, err := env.store.PutFile(ctx, key, bytes.NewReader(content), PutFileOptions{}) + if err != nil { + errors <- fmt.Errorf("goroutine %d, upload %d: %w", goroutineID, j, err) + return + } + + reader, _, err := env.store.GetFile(ctx, key) + if err != nil { + errors <- fmt.Errorf("goroutine %d, download %d: %w", goroutineID, j, err) + return + } + + got, err := io.ReadAll(reader) + reader.Close() + if err != nil { + errors <- fmt.Errorf("goroutine %d, read %d: %w", goroutineID, j, err) + return + } + + if !bytes.Equal(got, content) { + errors <- fmt.Errorf("goroutine %d, mismatch %d", goroutineID, j) + return + } + } + }(i) + } + + wg.Wait() + close(errors) + + var errList []error + for err := range errors { + errList = append(errList, err) + } + + if len(errList) > 0 { + t.Fatalf("Concurrent file upload errors: %v", errList) + } + + t.Logf("Successfully completed %d concurrent file uploads", numGoroutines*numFilesPerGoroutine) + }) +} + +// ============================================================================= +// Multi-Station Simulation +// ============================================================================= + +func TestIntegration_MultiStationWorkflow(t *testing.T) { + // Use shared embedded server + port := getFreePort(t) + httpPort := getFreePort(t) + + serverCfg := config.LatticeEmbeddedNATSConfig{ + Port: port, + HTTPPort: httpPort, + StoreDir: t.TempDir(), + } + + server := lattice.NewEmbeddedServer(serverCfg) + if err := server.Start(); err != nil { + t.Fatalf("Failed to start embedded NATS server: %v", err) + } + defer server.Shutdown() + + // Create two stations (simulating distributed agents) + createStation := func(name string) (*lattice.Client, *Store, *HandoffManager) { + clientCfg := config.LatticeConfig{ + StationID: name, + StationName: name + " Station", + NATS: config.LatticeNATSConfig{URL: server.ClientURL()}, + } + + client, err := lattice.NewClient(clientCfg) + if err != nil { + t.Fatalf("Failed to create client for %s: %v", name, err) + } + + if err := client.Connect(); err != nil { + t.Fatalf("Failed to connect %s: %v", name, err) + } + + store, err := NewStore(client.Conn(), DefaultStoreConfig()) + if err != nil { + client.Close() + t.Fatalf("Failed to create store for %s: %v", name, err) + } + + return client, store, NewHandoffManager(store) + } + + // Station A: Analyzer + clientA, storeA, handoffA := createStation("station-a") + defer storeA.Close() + defer clientA.Close() + + // Station B: Fixer + clientB, storeB, handoffB := createStation("station-b") + defer storeB.Close() + defer clientB.Close() + + ctx := context.Background() + + t.Run("CrossStationWorkflowHandoff", func(t *testing.T) { + workflowRunID := "wfr-cross-station" + + _, err := handoffA.StartWorkflow(ctx, StartWorkflowInput{ + WorkflowID: "wf-distributed", + WorkflowRunID: workflowRunID, + GitBranch: "agent/distributed-fix", + SharedData: map[string]interface{}{ + "initiated_by": "station-a", + }, + }) + if err != nil { + t.Fatalf("Station A StartWorkflow failed: %v", err) + } + + // Station A executes first step + runStateA, _, err := handoffA.StartStep(ctx, StartStepInput{ + WorkflowRunID: workflowRunID, + StepName: "analyze", + AgentName: "analyzer-a", + RunID: "run-a-1", + Task: "Analyze codebase", + }) + if err != nil { + t.Fatalf("Station A StartStep failed: %v", err) + } + + // Station A completes step and uploads output file + outputContent := []byte(`{"bugs": ["null-pointer", "race-condition"], "severity": "high"}`) + _, err = storeA.PutRunFile(ctx, runStateA.RunID, "analysis-report.json", bytes.NewReader(outputContent), PutFileOptions{ + ContentType: "application/json", + }) + if err != nil { + t.Fatalf("Station A PutRunFile failed: %v", err) + } + + err = handoffA.CompleteStep(ctx, CompleteStepInput{ + RunID: runStateA.RunID, + WorkflowRunID: workflowRunID, + Status: "completed", + Result: "Found 2 critical bugs", + Summary: "Analysis complete", + FilesModified: []string{}, + Commits: []string{}, + }) + if err != nil { + t.Fatalf("Station A CompleteStep failed: %v", err) + } + + t.Log("Station A completed analysis step") + + // Station B picks up the workflow (simulating handoff) + // Station B can read the workflow context because it's stored in shared NATS KV + prevCtx, err := handoffB.GetPreviousStepContext(ctx, workflowRunID) + if err != nil { + t.Fatalf("Station B GetPreviousStepContext failed: %v", err) + } + + if prevCtx == nil { + t.Fatal("Station B could not read previous step context from Station A") + } + + if prevCtx.StepName != "analyze" { + t.Errorf("Station B saw wrong previous step: got %q", prevCtx.StepName) + } + + t.Logf("Station B received handoff: previous step=%q, summary=%q", prevCtx.StepName, prevCtx.Summary) + + // Station B downloads output from Station A + reader, _, err := storeB.GetRunFile(ctx, runStateA.RunID, "analysis-report.json") + if err != nil { + t.Fatalf("Station B GetRunFile failed: %v", err) + } + + downloadedContent, err := io.ReadAll(reader) + reader.Close() + if err != nil { + t.Fatalf("Station B ReadAll failed: %v", err) + } + + if !bytes.Equal(downloadedContent, outputContent) { + t.Errorf("Station B received wrong content: got %q, want %q", downloadedContent, outputContent) + } + + t.Log("Station B downloaded analysis report from Station A") + + // Station B executes fix step + runStateB, _, err := handoffB.StartStep(ctx, StartStepInput{ + WorkflowRunID: workflowRunID, + StepName: "fix", + AgentName: "fixer-b", + RunID: "run-b-1", + Task: "Fix the 2 bugs found by analyzer", + }) + if err != nil { + t.Fatalf("Station B StartStep failed: %v", err) + } + + err = handoffB.CompleteStep(ctx, CompleteStepInput{ + RunID: runStateB.RunID, + WorkflowRunID: workflowRunID, + Status: "completed", + Result: "All bugs fixed", + Summary: "Fixes applied", + FilesModified: []string{"main.go", "handler.go"}, + Commits: []string{"fix123", "fix456"}, + }) + if err != nil { + t.Fatalf("Station B CompleteStep failed: %v", err) + } + + t.Log("Station B completed fix step") + + // Verify final state (readable from either station) + finalCtx, err := storeA.GetWorkflowContext(ctx, workflowRunID) + if err != nil { + t.Fatalf("Final GetWorkflowContext failed: %v", err) + } + + if len(finalCtx.Steps) != 2 { + t.Fatalf("Expected 2 steps, got %d", len(finalCtx.Steps)) + } + + if finalCtx.Steps[0].AgentName != "analyzer-a" { + t.Errorf("First step agent mismatch: got %q", finalCtx.Steps[0].AgentName) + } + if finalCtx.Steps[1].AgentName != "fixer-b" { + t.Errorf("Second step agent mismatch: got %q", finalCtx.Steps[1].AgentName) + } + + t.Logf("Cross-station workflow completed: %d steps across 2 stations", len(finalCtx.Steps)) + t.Logf(" Step 1: %s by %s on %s", finalCtx.Steps[0].StepName, finalCtx.Steps[0].AgentName, "station-a") + t.Logf(" Step 2: %s by %s on %s", finalCtx.Steps[1].StepName, finalCtx.Steps[1].AgentName, "station-b") + }) +} diff --git a/pkg/harness/nats/store.go b/pkg/harness/nats/store.go new file mode 100644 index 00000000..dd274c47 --- /dev/null +++ b/pkg/harness/nats/store.go @@ -0,0 +1,609 @@ +// Package nats provides NATS JetStream integration for the agentic harness. +// It handles inter-agent file sharing, workflow state, and context preservation. +package nats + +import ( + "context" + "encoding/json" + "fmt" + "io" + "strings" + "sync" + "time" + + "github.com/nats-io/nats.go" + "github.com/nats-io/nats.go/jetstream" +) + +const ( + DefaultKVBucket = "harness-state" + DefaultObjectBucket = "harness-files" + DefaultStateTTL = 24 * time.Hour + DefaultObjectStoreMaxGB = 1 + + PrefixRun = "run" + PrefixWorkflow = "workflow" + PrefixShared = "shared" +) + +// StoreConfig configures the NATS store. +type StoreConfig struct { + // KVBucket is the name of the KV bucket for state + KVBucket string `yaml:"kv_bucket" json:"kv_bucket"` + + // ObjectBucket is the name of the object store for files + ObjectBucket string `yaml:"object_bucket" json:"object_bucket"` + + // MaxFileSize is the maximum file size in bytes (default: 100MB) + MaxFileSize int64 `yaml:"max_file_size" json:"max_file_size"` + + // TTL is the default time-to-live for state entries + TTL time.Duration `yaml:"ttl" json:"ttl"` +} + +// DefaultStoreConfig returns sensible defaults. +func DefaultStoreConfig() StoreConfig { + return StoreConfig{ + KVBucket: DefaultKVBucket, + ObjectBucket: DefaultObjectBucket, + MaxFileSize: 100 * 1024 * 1024, // 100MB + TTL: DefaultStateTTL, + } +} + +// Store provides NATS-based storage for the agentic harness. +type Store struct { + js jetstream.JetStream + config StoreConfig + + mu sync.RWMutex + kv jetstream.KeyValue + objectStore nats.ObjectStore + + // Legacy JS context for object store (nats.go hasn't migrated object store to jetstream pkg yet) + legacyJS nats.JetStreamContext +} + +// NewStore creates a new NATS store with the given JetStream context. +func NewStore(nc *nats.Conn, config StoreConfig) (*Store, error) { + js, err := jetstream.New(nc) + if err != nil { + return nil, fmt.Errorf("create jetstream context: %w", err) + } + + // Get legacy JS context for object store operations + legacyJS, err := nc.JetStream() + if err != nil { + return nil, fmt.Errorf("create legacy jetstream context: %w", err) + } + + if config.KVBucket == "" { + config.KVBucket = DefaultKVBucket + } + if config.ObjectBucket == "" { + config.ObjectBucket = DefaultObjectBucket + } + if config.TTL == 0 { + config.TTL = DefaultStateTTL + } + if config.MaxFileSize == 0 { + config.MaxFileSize = 100 * 1024 * 1024 + } + + return &Store{ + js: js, + legacyJS: legacyJS, + config: config, + }, nil +} + +// getKV returns the KV bucket, creating it if necessary. +func (s *Store) getKV(ctx context.Context) (jetstream.KeyValue, error) { + s.mu.RLock() + if s.kv != nil { + kv := s.kv + s.mu.RUnlock() + return kv, nil + } + s.mu.RUnlock() + + s.mu.Lock() + defer s.mu.Unlock() + + // Double-check after acquiring write lock + if s.kv != nil { + return s.kv, nil + } + + // Try to get existing bucket + kv, err := s.js.KeyValue(ctx, s.config.KVBucket) + if err == nil { + s.kv = kv + return kv, nil + } + + // Create bucket if it doesn't exist + kv, err = s.js.CreateKeyValue(ctx, jetstream.KeyValueConfig{ + Bucket: s.config.KVBucket, + Description: "Agentic harness state storage", + TTL: s.config.TTL, + }) + if err != nil { + return nil, fmt.Errorf("create KV bucket %s: %w", s.config.KVBucket, err) + } + + s.kv = kv + return kv, nil +} + +// getObjectStore returns the object store bucket, creating it if necessary. +func (s *Store) getObjectStore(ctx context.Context) (nats.ObjectStore, error) { + s.mu.RLock() + if s.objectStore != nil { + os := s.objectStore + s.mu.RUnlock() + return os, nil + } + s.mu.RUnlock() + + s.mu.Lock() + defer s.mu.Unlock() + + // Double-check after acquiring write lock + if s.objectStore != nil { + return s.objectStore, nil + } + + // Try to get existing bucket + os, err := s.legacyJS.ObjectStore(s.config.ObjectBucket) + if err == nil { + s.objectStore = os + return os, nil + } + + // Create bucket if it doesn't exist + if err == nats.ErrBucketNotFound || err == nats.ErrStreamNotFound || strings.Contains(err.Error(), "stream not found") { + os, err = s.legacyJS.CreateObjectStore(&nats.ObjectStoreConfig{ + Bucket: s.config.ObjectBucket, + Description: "Agentic harness file storage", + MaxBytes: DefaultObjectStoreMaxGB * 1024 * 1024 * 1024, + }) + if err != nil { + return nil, fmt.Errorf("create object store %s: %w", s.config.ObjectBucket, err) + } + s.objectStore = os + return os, nil + } + + return nil, fmt.Errorf("get object store %s: %w", s.config.ObjectBucket, err) +} + +// ==================== KV State Operations ==================== + +// GetState retrieves a state value by key. +func (s *Store) GetState(ctx context.Context, key string) ([]byte, error) { + kv, err := s.getKV(ctx) + if err != nil { + return nil, err + } + + entry, err := kv.Get(ctx, key) + if err != nil { + if err == jetstream.ErrKeyNotFound { + return nil, nil + } + return nil, fmt.Errorf("get state %s: %w", key, err) + } + + return entry.Value(), nil +} + +// SetState stores a state value. +func (s *Store) SetState(ctx context.Context, key string, value []byte) error { + kv, err := s.getKV(ctx) + if err != nil { + return err + } + + if _, err := kv.Put(ctx, key, value); err != nil { + return fmt.Errorf("put state %s: %w", key, err) + } + + return nil +} + +// DeleteState removes a state value. +func (s *Store) DeleteState(ctx context.Context, key string) error { + kv, err := s.getKV(ctx) + if err != nil { + return err + } + + if err := kv.Delete(ctx, key); err != nil && err != jetstream.ErrKeyNotFound { + return fmt.Errorf("delete state %s: %w", key, err) + } + + return nil +} + +// GetJSON retrieves and unmarshals a JSON state value. +func (s *Store) GetJSON(ctx context.Context, key string, v interface{}) error { + data, err := s.GetState(ctx, key) + if err != nil { + return err + } + if data == nil { + return nil // Key not found, leave v as zero value + } + return json.Unmarshal(data, v) +} + +// SetJSON marshals and stores a JSON state value. +func (s *Store) SetJSON(ctx context.Context, key string, v interface{}) error { + data, err := json.Marshal(v) + if err != nil { + return fmt.Errorf("marshal state: %w", err) + } + return s.SetState(ctx, key, data) +} + +// ==================== Run State Operations ==================== + +// RunState represents the state of a single agent run. +type RunState struct { + RunID string `json:"run_id"` + AgentID string `json:"agent_id"` + AgentName string `json:"agent_name"` + WorkflowID string `json:"workflow_id,omitempty"` + StepName string `json:"step_name,omitempty"` + Status string `json:"status"` // running, completed, failed + StartedAt time.Time `json:"started_at"` + CompletedAt *time.Time `json:"completed_at,omitempty"` + Task string `json:"task"` + Result string `json:"result,omitempty"` + Error string `json:"error,omitempty"` + GitBranch string `json:"git_branch,omitempty"` + GitCommit string `json:"git_commit,omitempty"` + Metadata map[string]string `json:"metadata,omitempty"` +} + +// runStateKey returns the KV key for run state. +func runStateKey(runID string) string { + return fmt.Sprintf("%s.%s.state", PrefixRun, runID) +} + +// GetRunState retrieves the state of a run. +func (s *Store) GetRunState(ctx context.Context, runID string) (*RunState, error) { + var state RunState + if err := s.GetJSON(ctx, runStateKey(runID), &state); err != nil { + return nil, err + } + if state.RunID == "" { + return nil, nil // Not found + } + return &state, nil +} + +// SetRunState saves the state of a run. +func (s *Store) SetRunState(ctx context.Context, state *RunState) error { + return s.SetJSON(ctx, runStateKey(state.RunID), state) +} + +// ==================== Workflow Context Operations ==================== + +// WorkflowContext represents shared context across workflow steps. +type WorkflowContext struct { + WorkflowID string `json:"workflow_id"` + WorkflowRunID string `json:"workflow_run_id"` + StartedAt time.Time `json:"started_at"` + GitBranch string `json:"git_branch,omitempty"` + Steps []WorkflowStepSummary `json:"steps"` + SharedData map[string]interface{} `json:"shared_data,omitempty"` +} + +// WorkflowStepSummary represents a completed workflow step. +type WorkflowStepSummary struct { + StepName string `json:"step_name"` + AgentName string `json:"agent_name"` + RunID string `json:"run_id"` + Status string `json:"status"` + StartedAt time.Time `json:"started_at"` + CompletedAt *time.Time `json:"completed_at,omitempty"` + Summary string `json:"summary,omitempty"` + FilesModified []string `json:"files_modified,omitempty"` + Commits []string `json:"commits,omitempty"` +} + +// workflowContextKey returns the KV key for workflow context. +func workflowContextKey(workflowRunID string) string { + return fmt.Sprintf("%s.%s.context", PrefixWorkflow, workflowRunID) +} + +// GetWorkflowContext retrieves the context of a workflow run. +func (s *Store) GetWorkflowContext(ctx context.Context, workflowRunID string) (*WorkflowContext, error) { + var wctx WorkflowContext + if err := s.GetJSON(ctx, workflowContextKey(workflowRunID), &wctx); err != nil { + return nil, err + } + if wctx.WorkflowRunID == "" { + return nil, nil // Not found + } + return &wctx, nil +} + +// SetWorkflowContext saves the context of a workflow run. +func (s *Store) SetWorkflowContext(ctx context.Context, wctx *WorkflowContext) error { + return s.SetJSON(ctx, workflowContextKey(wctx.WorkflowRunID), wctx) +} + +// AddWorkflowStep adds a completed step to the workflow context. +func (s *Store) AddWorkflowStep(ctx context.Context, workflowRunID string, step WorkflowStepSummary) error { + wctx, err := s.GetWorkflowContext(ctx, workflowRunID) + if err != nil { + return err + } + if wctx == nil { + return fmt.Errorf("workflow context not found: %s", workflowRunID) + } + + wctx.Steps = append(wctx.Steps, step) + return s.SetWorkflowContext(ctx, wctx) +} + +// ==================== File Operations ==================== + +// FileMetadata contains metadata about a stored file. +type FileMetadata struct { + Key string `json:"key"` + Size int64 `json:"size"` + ContentType string `json:"content_type,omitempty"` + Checksum string `json:"checksum,omitempty"` + CreatedAt time.Time `json:"created_at"` + ExpiresAt *time.Time `json:"expires_at,omitempty"` + Metadata map[string]string `json:"metadata,omitempty"` +} + +// PutFileOptions configures file upload. +type PutFileOptions struct { + ContentType string + Description string + TTL time.Duration + Metadata map[string]string +} + +// PutFile stores a file in the object store. +func (s *Store) PutFile(ctx context.Context, key string, reader io.Reader, opts PutFileOptions) (*FileMetadata, error) { + os, err := s.getObjectStore(ctx) + if err != nil { + return nil, err + } + + // Read all data (object store requires this) + data, err := io.ReadAll(reader) + if err != nil { + return nil, fmt.Errorf("read file data: %w", err) + } + + if s.config.MaxFileSize > 0 && int64(len(data)) > s.config.MaxFileSize { + return nil, fmt.Errorf("file too large: %d bytes (max: %d)", len(data), s.config.MaxFileSize) + } + + meta := &nats.ObjectMeta{ + Name: key, + Description: opts.Description, + Headers: make(nats.Header), + } + + if opts.ContentType != "" { + meta.Headers.Set("Content-Type", opts.ContentType) + } + + if opts.TTL > 0 { + expiresAt := time.Now().Add(opts.TTL) + meta.Headers.Set("X-Expires-At", expiresAt.Format(time.RFC3339)) + } + + for k, v := range opts.Metadata { + meta.Headers.Set("X-Meta-"+k, v) + } + + _, err = os.Put(meta, strings.NewReader(string(data))) + if err != nil { + return nil, fmt.Errorf("put file %s: %w", key, err) + } + + info := &FileMetadata{ + Key: key, + Size: int64(len(data)), + ContentType: opts.ContentType, + CreatedAt: time.Now(), + Metadata: opts.Metadata, + } + + if opts.TTL > 0 { + expires := time.Now().Add(opts.TTL) + info.ExpiresAt = &expires + } + + return info, nil +} + +// GetFile retrieves a file from the object store. +func (s *Store) GetFile(ctx context.Context, key string) (io.ReadCloser, *FileMetadata, error) { + os, err := s.getObjectStore(ctx) + if err != nil { + return nil, nil, err + } + + result, err := os.Get(key) + if err != nil { + if err == nats.ErrObjectNotFound { + return nil, nil, nil // Not found + } + return nil, nil, fmt.Errorf("get file %s: %w", key, err) + } + + objInfo, err := result.Info() + if err != nil { + result.Close() + return nil, nil, fmt.Errorf("get file info %s: %w", key, err) + } + + info := &FileMetadata{ + Key: objInfo.Name, + Size: int64(objInfo.Size), + CreatedAt: objInfo.ModTime, + Metadata: make(map[string]string), + } + + if objInfo.Headers != nil { + info.ContentType = objInfo.Headers.Get("Content-Type") + info.Checksum = objInfo.Headers.Get("X-Checksum") + + if expiresStr := objInfo.Headers.Get("X-Expires-At"); expiresStr != "" { + if t, err := time.Parse(time.RFC3339, expiresStr); err == nil { + info.ExpiresAt = &t + } + } + + for k, values := range objInfo.Headers { + if strings.HasPrefix(k, "X-Meta-") && len(values) > 0 { + metaKey := strings.TrimPrefix(k, "X-Meta-") + info.Metadata[metaKey] = values[0] + } + } + } + + return result, info, nil +} + +// DeleteFile removes a file from the object store. +func (s *Store) DeleteFile(ctx context.Context, key string) error { + os, err := s.getObjectStore(ctx) + if err != nil { + return err + } + + if err := os.Delete(key); err != nil && err != nats.ErrObjectNotFound { + return fmt.Errorf("delete file %s: %w", key, err) + } + + return nil +} + +// ListFiles lists files with the given prefix. +func (s *Store) ListFiles(ctx context.Context, prefix string) ([]*FileMetadata, error) { + os, err := s.getObjectStore(ctx) + if err != nil { + return nil, err + } + + objects, err := os.List() + if err != nil { + return nil, fmt.Errorf("list files: %w", err) + } + + var files []*FileMetadata + for _, obj := range objects { + if prefix == "" || strings.HasPrefix(obj.Name, prefix) { + info := &FileMetadata{ + Key: obj.Name, + Size: int64(obj.Size), + CreatedAt: obj.ModTime, + Metadata: make(map[string]string), + } + + if obj.Headers != nil { + info.ContentType = obj.Headers.Get("Content-Type") + + for k, values := range obj.Headers { + if strings.HasPrefix(k, "X-Meta-") && len(values) > 0 { + metaKey := strings.TrimPrefix(k, "X-Meta-") + info.Metadata[metaKey] = values[0] + } + } + } + + files = append(files, info) + } + } + + return files, nil +} + +// ==================== Run File Operations ==================== + +// runFileKey returns the object store key for a run file. +func runFileKey(runID, filename string) string { + return fmt.Sprintf("%s/%s/output/%s", PrefixRun, runID, filename) +} + +// PutRunFile stores an output file for a run. +func (s *Store) PutRunFile(ctx context.Context, runID, filename string, reader io.Reader, opts PutFileOptions) (*FileMetadata, error) { + return s.PutFile(ctx, runFileKey(runID, filename), reader, opts) +} + +// GetRunFile retrieves an output file from a run. +func (s *Store) GetRunFile(ctx context.Context, runID, filename string) (io.ReadCloser, *FileMetadata, error) { + return s.GetFile(ctx, runFileKey(runID, filename)) +} + +// ListRunFiles lists output files from a run. +func (s *Store) ListRunFiles(ctx context.Context, runID string) ([]*FileMetadata, error) { + prefix := fmt.Sprintf("%s/%s/output/", PrefixRun, runID) + return s.ListFiles(ctx, prefix) +} + +// ==================== Shared File Operations ==================== + +// sharedFileKey returns the object store key for a shared file. +func sharedFileKey(key string) string { + return fmt.Sprintf("%s/%s", PrefixShared, key) +} + +// PutSharedFile stores a shared file accessible by any agent. +func (s *Store) PutSharedFile(ctx context.Context, key string, reader io.Reader, opts PutFileOptions) (*FileMetadata, error) { + return s.PutFile(ctx, sharedFileKey(key), reader, opts) +} + +// GetSharedFile retrieves a shared file. +func (s *Store) GetSharedFile(ctx context.Context, key string) (io.ReadCloser, *FileMetadata, error) { + return s.GetFile(ctx, sharedFileKey(key)) +} + +// ==================== Cleanup ==================== + +// CleanupExpiredFiles removes files that have expired. +func (s *Store) CleanupExpiredFiles(ctx context.Context) (int, error) { + os, err := s.getObjectStore(ctx) + if err != nil { + return 0, err + } + + objects, err := os.List() + if err != nil { + return 0, fmt.Errorf("list files: %w", err) + } + + now := time.Now() + var deleted int + + for _, obj := range objects { + if obj.Headers != nil { + if expiresStr := obj.Headers.Get("X-Expires-At"); expiresStr != "" { + if t, err := time.Parse(time.RFC3339, expiresStr); err == nil && now.After(t) { + if err := os.Delete(obj.Name); err == nil { + deleted++ + } + } + } + } + } + + return deleted, nil +} + +// Close releases resources (no-op for NATS stores). +func (s *Store) Close() error { + return nil +} diff --git a/pkg/harness/nats/store_test.go b/pkg/harness/nats/store_test.go new file mode 100644 index 00000000..d9eed42d --- /dev/null +++ b/pkg/harness/nats/store_test.go @@ -0,0 +1,168 @@ +package nats + +import ( + "testing" + "time" +) + +func TestDefaultStoreConfig(t *testing.T) { + config := DefaultStoreConfig() + + if config.KVBucket != DefaultKVBucket { + t.Errorf("expected KVBucket %q, got %q", DefaultKVBucket, config.KVBucket) + } + + if config.ObjectBucket != DefaultObjectBucket { + t.Errorf("expected ObjectBucket %q, got %q", DefaultObjectBucket, config.ObjectBucket) + } + + if config.MaxFileSize != 100*1024*1024 { + t.Errorf("expected MaxFileSize 100MB, got %d", config.MaxFileSize) + } + + if config.TTL != DefaultStateTTL { + t.Errorf("expected TTL %v, got %v", DefaultStateTTL, config.TTL) + } +} + +func TestRunStateKey(t *testing.T) { + key := runStateKey("run-123") + expected := "run.run-123.state" + if key != expected { + t.Errorf("expected %q, got %q", expected, key) + } +} + +func TestWorkflowContextKey(t *testing.T) { + key := workflowContextKey("wf-456") + expected := "workflow.wf-456.context" + if key != expected { + t.Errorf("expected %q, got %q", expected, key) + } +} + +func TestRunFileKey(t *testing.T) { + key := runFileKey("run-123", "output.txt") + expected := "run/run-123/output/output.txt" + if key != expected { + t.Errorf("expected %q, got %q", expected, key) + } +} + +func TestSharedFileKey(t *testing.T) { + key := sharedFileKey("myfile.json") + expected := "shared/myfile.json" + if key != expected { + t.Errorf("expected %q, got %q", expected, key) + } +} + +func TestRunState(t *testing.T) { + now := time.Now() + state := RunState{ + RunID: "run-123", + AgentID: "agent-1", + AgentName: "coder", + Status: "running", + StartedAt: now, + Task: "Fix the bug", + } + + if state.RunID != "run-123" { + t.Errorf("expected RunID %q, got %q", "run-123", state.RunID) + } + + if state.Status != "running" { + t.Errorf("expected Status %q, got %q", "running", state.Status) + } +} + +func TestWorkflowContext(t *testing.T) { + wctx := WorkflowContext{ + WorkflowID: "wf-1", + WorkflowRunID: "wfr-123", + StartedAt: time.Now(), + GitBranch: "feature/test", + Steps: []WorkflowStepSummary{}, + SharedData: map[string]interface{}{"key": "value"}, + } + + if wctx.WorkflowID != "wf-1" { + t.Errorf("expected WorkflowID %q, got %q", "wf-1", wctx.WorkflowID) + } + + if wctx.GitBranch != "feature/test" { + t.Errorf("expected GitBranch %q, got %q", "feature/test", wctx.GitBranch) + } + + if wctx.SharedData["key"] != "value" { + t.Errorf("expected SharedData key to be 'value'") + } +} + +func TestWorkflowStepSummary(t *testing.T) { + now := time.Now() + step := WorkflowStepSummary{ + StepName: "analyze", + AgentName: "analyzer", + RunID: "run-456", + Status: "completed", + StartedAt: now, + Summary: "Analyzed the code", + FilesModified: []string{"main.go", "util.go"}, + Commits: []string{"abc123"}, + } + + if step.StepName != "analyze" { + t.Errorf("expected StepName %q, got %q", "analyze", step.StepName) + } + + if len(step.FilesModified) != 2 { + t.Errorf("expected 2 files modified, got %d", len(step.FilesModified)) + } + + if len(step.Commits) != 1 { + t.Errorf("expected 1 commit, got %d", len(step.Commits)) + } +} + +func TestFileMetadata(t *testing.T) { + now := time.Now() + meta := FileMetadata{ + Key: "test/file.txt", + Size: 1024, + ContentType: "text/plain", + Checksum: "abc123", + CreatedAt: now, + Metadata: map[string]string{"custom": "value"}, + } + + if meta.Key != "test/file.txt" { + t.Errorf("expected Key %q, got %q", "test/file.txt", meta.Key) + } + + if meta.Size != 1024 { + t.Errorf("expected Size 1024, got %d", meta.Size) + } + + if meta.Metadata["custom"] != "value" { + t.Errorf("expected Metadata custom to be 'value'") + } +} + +func TestPutFileOptions(t *testing.T) { + opts := PutFileOptions{ + ContentType: "application/json", + Description: "Test file", + TTL: 24 * time.Hour, + Metadata: map[string]string{"source": "test"}, + } + + if opts.ContentType != "application/json" { + t.Errorf("expected ContentType %q, got %q", "application/json", opts.ContentType) + } + + if opts.TTL != 24*time.Hour { + t.Errorf("expected TTL 24h, got %v", opts.TTL) + } +} diff --git a/pkg/harness/prompt/builder.go b/pkg/harness/prompt/builder.go new file mode 100644 index 00000000..6fe7dae2 --- /dev/null +++ b/pkg/harness/prompt/builder.go @@ -0,0 +1,329 @@ +package prompt + +import ( + "fmt" + "os/exec" + "path/filepath" + "strings" + "time" +) + +type Builder struct { + agentPrompt string + workspacePath string + workspaceMode string + gitEnabled bool + gitBranch string + previousContext *PreviousStepContext + toolDescriptions []ToolDescription + customSections []CustomSection +} + +type ToolDescription struct { + Name string + Description string + Parameters string +} + +type PreviousStepContext struct { + StepName string + AgentName string + Summary string + FilesModified []string + Commits []string +} + +type CustomSection struct { + Title string + Content string +} + +func NewBuilder() *Builder { + return &Builder{} +} + +func (b *Builder) WithAgentPrompt(prompt string) *Builder { + b.agentPrompt = prompt + return b +} + +func (b *Builder) WithWorkspace(path, mode string) *Builder { + b.workspacePath = path + b.workspaceMode = mode + return b +} + +func (b *Builder) WithGit(enabled bool, branch string) *Builder { + b.gitEnabled = enabled + b.gitBranch = branch + return b +} + +func (b *Builder) WithPreviousContext(ctx *PreviousStepContext) *Builder { + b.previousContext = ctx + return b +} + +func (b *Builder) WithTools(tools []ToolDescription) *Builder { + b.toolDescriptions = tools + return b +} + +func (b *Builder) WithCustomSection(title, content string) *Builder { + b.customSections = append(b.customSections, CustomSection{ + Title: title, + Content: content, + }) + return b +} + +func (b *Builder) Build() string { + var sections []string + + sections = append(sections, b.buildEnvironmentSection()) + + if b.agentPrompt != "" { + sections = append(sections, b.buildAgentSection()) + } + + sections = append(sections, b.buildToolGuidelines()) + + if b.previousContext != nil { + sections = append(sections, b.buildWorkflowContext()) + } + + for _, custom := range b.customSections { + sections = append(sections, fmt.Sprintf("## %s\n\n%s", custom.Title, custom.Content)) + } + + return strings.Join(sections, "\n\n") +} + +func (b *Builder) buildEnvironmentSection() string { + var lines []string + lines = append(lines, "# Environment") + lines = append(lines, "") + + lines = append(lines, fmt.Sprintf("Working directory: %s", b.workspacePath)) + lines = append(lines, fmt.Sprintf("Workspace mode: %s", b.workspaceMode)) + lines = append(lines, fmt.Sprintf("Platform: %s", getPlatform())) + lines = append(lines, fmt.Sprintf("Date: %s", time.Now().Format("Mon Jan 02 2006"))) + + if b.gitEnabled { + lines = append(lines, fmt.Sprintf("Git enabled: yes")) + if b.gitBranch != "" { + lines = append(lines, fmt.Sprintf("Git branch: %s", b.gitBranch)) + } + if isGitRepo(b.workspacePath) { + lines = append(lines, "Is git repo: yes") + } + } + + return strings.Join(lines, "\n") +} + +func (b *Builder) buildAgentSection() string { + return fmt.Sprintf("# Agent Instructions\n\n%s", b.agentPrompt) +} + +func (b *Builder) buildToolGuidelines() string { + var lines []string + lines = append(lines, "# Tool Guidelines") + lines = append(lines, "") + + lines = append(lines, "## File Operations") + lines = append(lines, "- Use `read` to view file contents with line numbers") + lines = append(lines, "- Use `write` to create new files (auto-creates directories)") + lines = append(lines, "- Use `edit` for precise string replacements in existing files") + lines = append(lines, "- Use `glob` to find files by pattern (e.g., `**/*.go`)") + lines = append(lines, "- Use `grep` to search file contents with regex") + lines = append(lines, "") + + lines = append(lines, "## Bash Commands") + lines = append(lines, "- Always provide clear descriptions for commands") + lines = append(lines, "- Commands are executed in the workspace directory") + lines = append(lines, "- Avoid destructive commands (rm -rf, etc.)") + lines = append(lines, "- Use && to chain related commands") + lines = append(lines, "- Maximum timeout: 2 minutes") + lines = append(lines, "") + + if b.gitEnabled { + lines = append(lines, "## Git Operations") + lines = append(lines, "- Check git status before and after changes") + lines = append(lines, "- NEVER force push or hard reset") + lines = append(lines, "- Commit messages should be descriptive") + lines = append(lines, "- Push operations may require approval") + lines = append(lines, "") + } + + lines = append(lines, "## Best Practices") + lines = append(lines, "- Read files before editing to understand context") + lines = append(lines, "- Make small, incremental changes") + lines = append(lines, "- Verify changes compile/run before completing") + lines = append(lines, "- Explain your reasoning when making decisions") + + return strings.Join(lines, "\n") +} + +func (b *Builder) buildWorkflowContext() string { + var lines []string + lines = append(lines, "# Previous Step Context") + lines = append(lines, "") + + ctx := b.previousContext + lines = append(lines, fmt.Sprintf("Previous step: %s (agent: %s)", ctx.StepName, ctx.AgentName)) + lines = append(lines, "") + + if ctx.Summary != "" { + lines = append(lines, "Summary:") + lines = append(lines, ctx.Summary) + lines = append(lines, "") + } + + if len(ctx.FilesModified) > 0 { + lines = append(lines, "Files modified:") + for _, f := range ctx.FilesModified { + lines = append(lines, fmt.Sprintf(" - %s", f)) + } + lines = append(lines, "") + } + + if len(ctx.Commits) > 0 { + lines = append(lines, "Commits:") + for _, c := range ctx.Commits { + lines = append(lines, fmt.Sprintf(" - %s", c)) + } + lines = append(lines, "") + } + + lines = append(lines, "Continue from where the previous agent left off. The workspace and git branch are preserved.") + + return strings.Join(lines, "\n") +} + +func getPlatform() string { + return "linux" +} + +func isGitRepo(path string) bool { + gitPath := filepath.Join(path, ".git") + cmd := exec.Command("test", "-d", gitPath) + return cmd.Run() == nil +} + +var defaultToolDescriptions = []ToolDescription{ + { + Name: "read", + Description: "Read file contents with line numbers", + Parameters: "path (required), offset (optional), limit (optional)", + }, + { + Name: "write", + Description: "Create or overwrite a file", + Parameters: "path (required), content (required)", + }, + { + Name: "edit", + Description: "Replace exact string in file", + Parameters: "path (required), oldString (required), newString (required), replaceAll (optional)", + }, + { + Name: "bash", + Description: "Execute shell command", + Parameters: "command (required), description (required), workdir (optional), timeout (optional)", + }, + { + Name: "glob", + Description: "Find files matching pattern", + Parameters: "pattern (required), path (optional)", + }, + { + Name: "grep", + Description: "Search file contents with regex", + Parameters: "pattern (required), path (optional), include (optional)", + }, +} + +func DefaultToolDescriptions() []ToolDescription { + return defaultToolDescriptions +} + +type PromptTemplate struct { + Name string + Content string +} + +var templates = map[string]PromptTemplate{ + "coder": { + Name: "coder", + Content: `You are an expert software engineer with deep knowledge of multiple programming languages and best practices. + +Your role is to help implement, debug, and improve code. You have access to the workspace filesystem and can read, write, and edit files directly. + +Key principles: +1. Write clean, maintainable code following established patterns in the codebase +2. Add appropriate tests for new functionality +3. Handle errors gracefully +4. Document complex logic with clear comments +5. Follow the existing code style and conventions + +When making changes: +- First understand the existing code structure +- Make minimal, focused changes +- Verify changes work before completing +- Explain your reasoning for significant decisions`, + }, + "reviewer": { + Name: "reviewer", + Content: `You are an expert code reviewer focused on code quality, security, and best practices. + +Your role is to review code changes and provide constructive feedback. You have read access to the codebase to understand context. + +Review criteria: +1. Code correctness and logic errors +2. Security vulnerabilities +3. Performance issues +4. Code style and consistency +5. Test coverage +6. Documentation quality + +When reviewing: +- Be specific about issues and suggested fixes +- Explain the reasoning behind recommendations +- Prioritize critical issues over style preferences +- Acknowledge good patterns and practices`, + }, + "debugger": { + Name: "debugger", + Content: `You are an expert debugger skilled at identifying and fixing software issues. + +Your role is to diagnose problems, identify root causes, and implement fixes. You have full access to the codebase and can run commands to investigate. + +Debugging approach: +1. Understand the reported issue and expected behavior +2. Reproduce the problem if possible +3. Analyze relevant code paths +4. Identify potential root causes +5. Test hypotheses systematically +6. Implement and verify the fix + +When debugging: +- Start with the most likely causes +- Use logging and debugging tools effectively +- Consider edge cases and race conditions +- Document findings and fix rationale`, + }, +} + +func GetTemplate(name string) (PromptTemplate, bool) { + t, ok := templates[name] + return t, ok +} + +func ListTemplates() []string { + names := make([]string, 0, len(templates)) + for name := range templates { + names = append(names, name) + } + return names +} diff --git a/pkg/harness/prompt/builder_test.go b/pkg/harness/prompt/builder_test.go new file mode 100644 index 00000000..a2a7ec3d --- /dev/null +++ b/pkg/harness/prompt/builder_test.go @@ -0,0 +1,223 @@ +package prompt + +import ( + "strings" + "testing" +) + +func TestBuilder_Build_Minimal(t *testing.T) { + builder := NewBuilder(). + WithWorkspace("/workspace", "host") + + result := builder.Build() + + if !strings.Contains(result, "Working directory: /workspace") { + t.Error("expected working directory in output") + } + + if !strings.Contains(result, "Workspace mode: host") { + t.Error("expected workspace mode in output") + } + + if !strings.Contains(result, "# Tool Guidelines") { + t.Error("expected tool guidelines section") + } +} + +func TestBuilder_Build_WithAgentPrompt(t *testing.T) { + builder := NewBuilder(). + WithWorkspace("/workspace", "host"). + WithAgentPrompt("You are a helpful coding assistant.") + + result := builder.Build() + + if !strings.Contains(result, "# Agent Instructions") { + t.Error("expected agent instructions section") + } + + if !strings.Contains(result, "You are a helpful coding assistant.") { + t.Error("expected agent prompt in output") + } +} + +func TestBuilder_Build_WithGit(t *testing.T) { + builder := NewBuilder(). + WithWorkspace("/workspace", "host"). + WithGit(true, "feature/test") + + result := builder.Build() + + if !strings.Contains(result, "Git enabled: yes") { + t.Error("expected git enabled in output") + } + + if !strings.Contains(result, "Git branch: feature/test") { + t.Error("expected git branch in output") + } + + if !strings.Contains(result, "## Git Operations") { + t.Error("expected git operations section") + } +} + +func TestBuilder_Build_WithPreviousContext(t *testing.T) { + builder := NewBuilder(). + WithWorkspace("/workspace", "host"). + WithPreviousContext(&PreviousStepContext{ + StepName: "analyze", + AgentName: "analyzer", + Summary: "Analyzed the codebase", + FilesModified: []string{"main.go", "util.go"}, + Commits: []string{"abc123", "def456"}, + }) + + result := builder.Build() + + if !strings.Contains(result, "# Previous Step Context") { + t.Error("expected previous step context section") + } + + if !strings.Contains(result, "Previous step: analyze (agent: analyzer)") { + t.Error("expected step and agent info") + } + + if !strings.Contains(result, "Analyzed the codebase") { + t.Error("expected summary in output") + } + + if !strings.Contains(result, "main.go") { + t.Error("expected files modified in output") + } + + if !strings.Contains(result, "abc123") { + t.Error("expected commits in output") + } +} + +func TestBuilder_Build_WithCustomSection(t *testing.T) { + builder := NewBuilder(). + WithWorkspace("/workspace", "host"). + WithCustomSection("Project Rules", "Always use gofmt before committing.") + + result := builder.Build() + + if !strings.Contains(result, "## Project Rules") { + t.Error("expected custom section title") + } + + if !strings.Contains(result, "Always use gofmt before committing.") { + t.Error("expected custom section content") + } +} + +func TestBuilder_Build_Complete(t *testing.T) { + builder := NewBuilder(). + WithWorkspace("/workspace", "host"). + WithAgentPrompt("You are an expert coder."). + WithGit(true, "main"). + WithPreviousContext(&PreviousStepContext{ + StepName: "setup", + AgentName: "initializer", + Summary: "Set up the project", + }). + WithCustomSection("Custom", "Custom content") + + result := builder.Build() + + sections := []string{ + "# Environment", + "# Agent Instructions", + "# Tool Guidelines", + "# Previous Step Context", + "## Custom", + } + + for _, section := range sections { + if !strings.Contains(result, section) { + t.Errorf("expected section %q in output", section) + } + } +} + +func TestGetTemplate(t *testing.T) { + template, ok := GetTemplate("coder") + if !ok { + t.Fatal("expected to find coder template") + } + + if template.Name != "coder" { + t.Errorf("expected template name 'coder', got %q", template.Name) + } + + if !strings.Contains(template.Content, "expert software engineer") { + t.Error("expected coder template to mention software engineer") + } +} + +func TestGetTemplate_NotFound(t *testing.T) { + _, ok := GetTemplate("nonexistent") + if ok { + t.Error("expected template not to be found") + } +} + +func TestListTemplates(t *testing.T) { + names := ListTemplates() + + if len(names) == 0 { + t.Error("expected at least one template") + } + + found := false + for _, name := range names { + if name == "coder" { + found = true + break + } + } + + if !found { + t.Error("expected to find 'coder' in template list") + } +} + +func TestDefaultToolDescriptions(t *testing.T) { + tools := DefaultToolDescriptions() + + if len(tools) == 0 { + t.Error("expected at least one tool description") + } + + foundBash := false + for _, tool := range tools { + if tool.Name == "bash" { + foundBash = true + if tool.Description == "" { + t.Error("expected bash tool to have description") + } + break + } + } + + if !foundBash { + t.Error("expected to find 'bash' in tool descriptions") + } +} + +func TestBuilder_WithTools(t *testing.T) { + tools := []ToolDescription{ + {Name: "custom", Description: "Custom tool", Parameters: "arg1, arg2"}, + } + + builder := NewBuilder(). + WithWorkspace("/workspace", "host"). + WithTools(tools) + + if len(builder.toolDescriptions) != 1 { + t.Errorf("expected 1 tool description, got %d", len(builder.toolDescriptions)) + } + + if builder.toolDescriptions[0].Name != "custom" { + t.Error("expected custom tool") + } +} diff --git a/pkg/harness/prompt/loader.go b/pkg/harness/prompt/loader.go new file mode 100644 index 00000000..9f6db53e --- /dev/null +++ b/pkg/harness/prompt/loader.go @@ -0,0 +1,236 @@ +package prompt + +import ( + "fmt" + "os" + "strings" + "time" + + "gopkg.in/yaml.v3" + + "station/pkg/harness/memory" + "station/pkg/harness/skills" +) + +// AgentConfig represents a parsed agent dotprompt with extended harness configuration +type AgentConfig struct { + // Standard Genkit/Station dotprompt fields + Model string `yaml:"model"` + Metadata AgentMetadata `yaml:"metadata"` + Tools []string `yaml:"tools,omitempty"` + Sandbox *SandboxConfig `yaml:"sandbox,omitempty"` + + // Extended harness configuration + Harness *HarnessConfig `yaml:"harness,omitempty"` + Workspace *WorkspaceConfig `yaml:"workspace,omitempty"` + Skills *skills.SkillsConfig `yaml:"skills,omitempty"` + Memory *memory.MemoryConfig `yaml:"memory,omitempty"` + + // The prompt content (system + user messages) + PromptContent string `yaml:"-"` +} + +// AgentMetadata contains agent metadata +type AgentMetadata struct { + Name string `yaml:"name"` + Description string `yaml:"description,omitempty"` + Tags []string `yaml:"tags,omitempty"` + MaxSteps int `yaml:"max_steps,omitempty"` +} + +// HarnessConfig configures the agentic harness behavior +type HarnessConfig struct { + Enabled bool `yaml:"enabled,omitempty"` + MaxSteps int `yaml:"max_steps,omitempty"` + Timeout string `yaml:"timeout,omitempty"` // e.g., "30m", "1h" + DoomLoop *DoomLoopConfig `yaml:"doom_loop,omitempty"` + Compaction *CompactionConfig `yaml:"compaction,omitempty"` + Progress *ProgressConfig `yaml:"progress,omitempty"` +} + +// DoomLoopConfig configures doom loop detection +type DoomLoopConfig struct { + Threshold int `yaml:"threshold,omitempty"` // Default: 3 + Action string `yaml:"action,omitempty"` // "summarize", "pause", "fail" +} + +// CompactionConfig configures context compaction +type CompactionConfig struct { + Enabled bool `yaml:"enabled,omitempty"` + Threshold float64 `yaml:"threshold,omitempty"` // Default: 0.85 + ProtectTokens int `yaml:"protect_tokens,omitempty"` // Default: 40000 + HistoryOffload bool `yaml:"history_offload,omitempty"` // Save to NATS before summarizing +} + +// ProgressConfig configures progress tracking +type ProgressConfig struct { + File string `yaml:"file,omitempty"` // e.g., "progress.md" + SyncInterval string `yaml:"sync_interval,omitempty"` // e.g., "5s" +} + +// WorkspaceConfig configures the working directory +type WorkspaceConfig struct { + Mode string `yaml:"mode,omitempty"` // "host" or "sandbox" + Scope string `yaml:"scope,omitempty"` // "session" or "workflow" + Cleanup string `yaml:"cleanup,omitempty"` // "always", "on_success", "never" +} + +// SandboxConfig configures sandbox isolation +type SandboxConfig struct { + Mode string `yaml:"mode,omitempty"` // "host", "docker", "e2b", "modal", "runloop", "daytona" + Image string `yaml:"image,omitempty"` + Timeout string `yaml:"timeout,omitempty"` + Resources *ResourceConfig `yaml:"resources,omitempty"` + Network *NetworkConfig `yaml:"network,omitempty"` +} + +// ResourceConfig defines resource limits +type ResourceConfig struct { + CPU string `yaml:"cpu,omitempty"` // e.g., "2" + Memory string `yaml:"memory,omitempty"` // e.g., "4g" +} + +// NetworkConfig defines network access +type NetworkConfig struct { + Enabled bool `yaml:"enabled,omitempty"` + AllowedHosts []string `yaml:"allowed_hosts,omitempty"` +} + +// LoadAgentConfig loads and parses an agent dotprompt file with extended harness fields +func LoadAgentConfig(path string) (*AgentConfig, error) { + content, err := os.ReadFile(path) + if err != nil { + return nil, fmt.Errorf("failed to read agent file: %w", err) + } + + return ParseAgentConfig(string(content)) +} + +// ParseAgentConfig parses agent config from dotprompt content string +func ParseAgentConfig(content string) (*AgentConfig, error) { + // Split frontmatter and prompt content + frontmatter, promptContent, err := splitFrontmatter(content) + if err != nil { + return nil, err + } + + // Parse frontmatter + var config AgentConfig + if err := yaml.Unmarshal([]byte(frontmatter), &config); err != nil { + return nil, fmt.Errorf("failed to parse frontmatter: %w", err) + } + + config.PromptContent = promptContent + return &config, nil +} + +// splitFrontmatter separates YAML frontmatter from prompt content +func splitFrontmatter(content string) (frontmatter, body string, err error) { + if !strings.HasPrefix(content, "---") { + return "", content, nil // No frontmatter + } + + // Find the closing --- + rest := content[3:] + parts := strings.SplitN(rest, "---", 2) + if len(parts) < 2 { + return "", "", fmt.Errorf("frontmatter not closed") + } + + frontmatter = strings.TrimSpace(parts[0]) + body = strings.TrimSpace(parts[1]) + return frontmatter, body, nil +} + +// IsAgenticHarness returns true if this agent uses the agentic harness +func (c *AgentConfig) IsAgenticHarness() bool { + return c.Harness != nil && c.Harness.Enabled +} + +// GetMaxSteps returns the effective max steps (harness > metadata > default) +func (c *AgentConfig) GetMaxSteps() int { + if c.Harness != nil && c.Harness.MaxSteps > 0 { + return c.Harness.MaxSteps + } + if c.Metadata.MaxSteps > 0 { + return c.Metadata.MaxSteps + } + return 50 // Default +} + +// GetTimeout returns the timeout as a Duration +func (c *AgentConfig) GetTimeout() time.Duration { + if c.Harness == nil || c.Harness.Timeout == "" { + return 30 * time.Minute // Default + } + + d, err := time.ParseDuration(c.Harness.Timeout) + if err != nil { + return 30 * time.Minute + } + return d +} + +// GetDoomLoopThreshold returns the doom loop threshold +func (c *AgentConfig) GetDoomLoopThreshold() int { + if c.Harness != nil && c.Harness.DoomLoop != nil && c.Harness.DoomLoop.Threshold > 0 { + return c.Harness.DoomLoop.Threshold + } + return 3 // Default +} + +// GetCompactionThreshold returns the compaction threshold +func (c *AgentConfig) GetCompactionThreshold() float64 { + if c.Harness != nil && c.Harness.Compaction != nil && c.Harness.Compaction.Threshold > 0 { + return c.Harness.Compaction.Threshold + } + return 0.85 // Default +} + +// GetProtectTokens returns the number of tokens to protect from compaction +func (c *AgentConfig) GetProtectTokens() int { + if c.Harness != nil && c.Harness.Compaction != nil && c.Harness.Compaction.ProtectTokens > 0 { + return c.Harness.Compaction.ProtectTokens + } + return 40000 // Default +} + +// IsCompactionEnabled returns whether compaction is enabled +func (c *AgentConfig) IsCompactionEnabled() bool { + if c.Harness == nil || c.Harness.Compaction == nil { + return true // Enabled by default for harness agents + } + return c.Harness.Compaction.Enabled +} + +// IsHistoryOffloadEnabled returns whether history offload is enabled +func (c *AgentConfig) IsHistoryOffloadEnabled() bool { + if c.Harness == nil || c.Harness.Compaction == nil { + return false // Disabled by default + } + return c.Harness.Compaction.HistoryOffload +} + +// GetSkillSources returns skill sources, using defaults if not specified +func (c *AgentConfig) GetSkillSources(envPath string) []string { + if c.Skills != nil && len(c.Skills.Sources) > 0 { + return c.Skills.Sources + } + return skills.DefaultSkillSources(envPath) +} + +// GetMemorySources returns memory sources, using defaults if not specified +func (c *AgentConfig) GetMemorySources(envPath string) []string { + if c.Memory != nil && len(c.Memory.Sources) > 0 { + return c.Memory.Sources + } + return memory.DefaultMemorySources(envPath) +} + +// GetSandboxMode returns the sandbox mode (defaults to "host") +func (c *AgentConfig) GetSandboxMode() string { + if c.Sandbox != nil && c.Sandbox.Mode != "" { + return c.Sandbox.Mode + } + return "host" +} diff --git a/pkg/harness/prompt/loader_test.go b/pkg/harness/prompt/loader_test.go new file mode 100644 index 00000000..0ab4dcd0 --- /dev/null +++ b/pkg/harness/prompt/loader_test.go @@ -0,0 +1,284 @@ +package prompt + +import ( + "testing" + "time" + + "station/pkg/harness/memory" + "station/pkg/harness/skills" +) + +func TestParseAgentConfig_Basic(t *testing.T) { + content := `--- +model: gpt-4o-mini +metadata: + name: test-agent + description: A test agent + max_steps: 25 +tools: + - bash + - read_file +--- + +You are a test agent.` + + config, err := ParseAgentConfig(content) + if err != nil { + t.Fatalf("ParseAgentConfig() error = %v", err) + } + + if config.Model != "gpt-4o-mini" { + t.Errorf("Model = %q, want gpt-4o-mini", config.Model) + } + if config.Metadata.Name != "test-agent" { + t.Errorf("Name = %q, want test-agent", config.Metadata.Name) + } + if len(config.Tools) != 2 { + t.Errorf("Tools count = %d, want 2", len(config.Tools)) + } + if config.PromptContent != "You are a test agent." { + t.Errorf("PromptContent = %q", config.PromptContent) + } +} + +func TestParseAgentConfig_WithHarness(t *testing.T) { + content := `--- +model: openai/gpt-4o +metadata: + name: harness-agent + description: An agentic harness agent +harness: + enabled: true + max_steps: 100 + timeout: 60m + doom_loop: + threshold: 5 + action: summarize + compaction: + enabled: true + threshold: 0.85 + protect_tokens: 50000 + history_offload: true +--- + +{{role "system"}} +You are a powerful agent.` + + config, err := ParseAgentConfig(content) + if err != nil { + t.Fatalf("ParseAgentConfig() error = %v", err) + } + + if !config.IsAgenticHarness() { + t.Error("should be agentic harness") + } + + if config.GetMaxSteps() != 100 { + t.Errorf("MaxSteps = %d, want 100", config.GetMaxSteps()) + } + + if config.GetTimeout() != 60*time.Minute { + t.Errorf("Timeout = %v, want 60m", config.GetTimeout()) + } + + if config.GetDoomLoopThreshold() != 5 { + t.Errorf("DoomLoopThreshold = %d, want 5", config.GetDoomLoopThreshold()) + } + + if !config.IsCompactionEnabled() { + t.Error("compaction should be enabled") + } + + if config.GetCompactionThreshold() != 0.85 { + t.Errorf("CompactionThreshold = %f, want 0.85", config.GetCompactionThreshold()) + } + + if config.GetProtectTokens() != 50000 { + t.Errorf("ProtectTokens = %d, want 50000", config.GetProtectTokens()) + } + + if !config.IsHistoryOffloadEnabled() { + t.Error("history offload should be enabled") + } +} + +func TestParseAgentConfig_WithSkillsAndMemory(t *testing.T) { + content := `--- +model: gpt-4o-mini +metadata: + name: skilled-agent +harness: + enabled: true +skills: + sources: + - /custom/skills + - .station/skills +memory: + sources: + - ~/.config/station/AGENTS.md + - .station/AGENTS.md +--- + +Agent prompt here.` + + config, err := ParseAgentConfig(content) + if err != nil { + t.Fatalf("ParseAgentConfig() error = %v", err) + } + + if len(config.Skills.Sources) != 2 { + t.Errorf("Skills sources = %d, want 2", len(config.Skills.Sources)) + } + + if len(config.Memory.Sources) != 2 { + t.Errorf("Memory sources = %d, want 2", len(config.Memory.Sources)) + } + + skillSources := config.GetSkillSources("/env") + if skillSources[0] != "/custom/skills" { + t.Errorf("First skill source = %q, want /custom/skills", skillSources[0]) + } +} + +func TestParseAgentConfig_WithSandbox(t *testing.T) { + content := `--- +model: gpt-4o-mini +metadata: + name: sandboxed-agent +sandbox: + mode: docker + image: python:3.12-slim + timeout: 30m + resources: + cpu: "2" + memory: 4g + network: + enabled: true + allowed_hosts: + - "*.github.com" + - "pypi.org" +--- + +Agent prompt.` + + config, err := ParseAgentConfig(content) + if err != nil { + t.Fatalf("ParseAgentConfig() error = %v", err) + } + + if config.GetSandboxMode() != "docker" { + t.Errorf("SandboxMode = %q, want docker", config.GetSandboxMode()) + } + + if config.Sandbox.Image != "python:3.12-slim" { + t.Errorf("Image = %q, want python:3.12-slim", config.Sandbox.Image) + } + + if config.Sandbox.Resources.CPU != "2" { + t.Errorf("CPU = %q, want 2", config.Sandbox.Resources.CPU) + } + + if !config.Sandbox.Network.Enabled { + t.Error("network should be enabled") + } + + if len(config.Sandbox.Network.AllowedHosts) != 2 { + t.Errorf("AllowedHosts count = %d, want 2", len(config.Sandbox.Network.AllowedHosts)) + } +} + +func TestParseAgentConfig_NoFrontmatter(t *testing.T) { + content := "Just a prompt without frontmatter." + + config, err := ParseAgentConfig(content) + if err != nil { + t.Fatalf("ParseAgentConfig() error = %v", err) + } + + if config.PromptContent != content { + t.Errorf("PromptContent = %q, want original content", config.PromptContent) + } +} + +func TestParseAgentConfig_UnclosedFrontmatter(t *testing.T) { + content := "---\nmodel: test" + + _, err := ParseAgentConfig(content) + if err == nil { + t.Error("should error on unclosed frontmatter") + } +} + +func TestAgentConfig_Defaults(t *testing.T) { + config := &AgentConfig{} + + if config.IsAgenticHarness() { + t.Error("should not be harness by default") + } + + if config.GetMaxSteps() != 50 { + t.Errorf("default MaxSteps = %d, want 50", config.GetMaxSteps()) + } + + if config.GetTimeout() != 30*time.Minute { + t.Errorf("default Timeout = %v, want 30m", config.GetTimeout()) + } + + if config.GetDoomLoopThreshold() != 3 { + t.Errorf("default DoomLoopThreshold = %d, want 3", config.GetDoomLoopThreshold()) + } + + if config.GetCompactionThreshold() != 0.85 { + t.Errorf("default CompactionThreshold = %f, want 0.85", config.GetCompactionThreshold()) + } + + if config.GetProtectTokens() != 40000 { + t.Errorf("default ProtectTokens = %d, want 40000", config.GetProtectTokens()) + } + + if config.GetSandboxMode() != "host" { + t.Errorf("default SandboxMode = %q, want host", config.GetSandboxMode()) + } +} + +func TestAgentConfig_GetSkillSources_CustomVsDefault(t *testing.T) { + // With custom sources + config := &AgentConfig{ + Skills: &skills.SkillsConfig{ + Sources: []string{"/custom/path"}, + }, + } + + sources := config.GetSkillSources("/env/default") + if len(sources) != 1 || sources[0] != "/custom/path" { + t.Error("should return custom sources when specified") + } + + // Without custom sources - should get defaults + config2 := &AgentConfig{} + sources2 := config2.GetSkillSources("/env/default") + if len(sources2) < 2 { + t.Error("should return default sources when not specified") + } +} + +func TestAgentConfig_GetMemorySources_CustomVsDefault(t *testing.T) { + // With custom sources + config := &AgentConfig{ + Memory: &memory.MemoryConfig{ + Sources: []string{"/custom/AGENTS.md"}, + }, + } + + sources := config.GetMemorySources("/env/default") + if len(sources) != 1 || sources[0] != "/custom/AGENTS.md" { + t.Error("should return custom sources when specified") + } + + // Without custom sources - should get defaults + config2 := &AgentConfig{} + sources2 := config2.GetMemorySources("/env/default") + if len(sources2) < 2 { + t.Error("should return default sources when not specified") + } +} diff --git a/pkg/harness/sandbox/docker.go b/pkg/harness/sandbox/docker.go new file mode 100644 index 00000000..65dedf9f --- /dev/null +++ b/pkg/harness/sandbox/docker.go @@ -0,0 +1,489 @@ +package sandbox + +import ( + "bytes" + "context" + "encoding/base64" + "encoding/json" + "fmt" + "io" + "os" + "os/exec" + "path/filepath" + "strings" + "time" + + "github.com/google/uuid" +) + +type DockerSandbox struct { + id string + containerID string + config Config + created bool +} + +func NewDockerSandbox(cfg Config) (*DockerSandbox, error) { + if cfg.Image == "" { + cfg.Image = "ubuntu:22.04" + } + + return &DockerSandbox{ + id: uuid.New().String(), + config: cfg, + }, nil +} + +func (s *DockerSandbox) Create(ctx context.Context) error { + if s.created { + return nil + } + + if err := s.pullImage(ctx); err != nil { + return fmt.Errorf("failed to pull image: %w", err) + } + + args := []string{ + "create", + "--name", s.containerName(), + "--rm", + } + + if s.config.Resources.CPU > 0 { + args = append(args, "--cpus", fmt.Sprintf("%.2f", s.config.Resources.CPU)) + } + + if s.config.Resources.Memory != "" { + args = append(args, "--memory", s.config.Resources.Memory) + } + + if s.config.Resources.PIDs > 0 { + args = append(args, "--pids-limit", fmt.Sprintf("%d", s.config.Resources.PIDs)) + } + + if !s.config.Network.Enabled { + args = append(args, "--network", "none") + } + + if s.config.WorkspacePath != "" { + absPath, err := filepath.Abs(s.config.WorkspacePath) + if err != nil { + return fmt.Errorf("failed to get absolute workspace path: %w", err) + } + + if err := os.MkdirAll(absPath, 0755); err != nil { + return fmt.Errorf("failed to create workspace directory: %w", err) + } + + args = append(args, "-v", fmt.Sprintf("%s:/workspace", absPath)) + args = append(args, "-w", "/workspace") + } + + for k, v := range s.config.Environment { + args = append(args, "-e", fmt.Sprintf("%s=%s", k, v)) + } + + args = append(args, "--cap-drop", "ALL") + args = append(args, "--security-opt", "no-new-privileges") + args = append(args, "--read-only") + args = append(args, "--tmpfs", "/tmp:rw,noexec,nosuid,size=100m") + + args = append(args, s.config.Image, "tail", "-f", "/dev/null") + + cmd := exec.CommandContext(ctx, "docker", args...) + output, err := cmd.CombinedOutput() + if err != nil { + return fmt.Errorf("failed to create container: %w, output: %s", err, string(output)) + } + + s.containerID = strings.TrimSpace(string(output)) + + startCmd := exec.CommandContext(ctx, "docker", "start", s.containerName()) + if output, err := startCmd.CombinedOutput(); err != nil { + return fmt.Errorf("failed to start container: %w, output: %s", err, string(output)) + } + + s.created = true + return nil +} + +func (s *DockerSandbox) Exec(ctx context.Context, command string, args ...string) (*ExecResult, error) { + return s.ExecWithStdin(ctx, nil, command, args...) +} + +func (s *DockerSandbox) ExecWithStdin(ctx context.Context, stdin io.Reader, command string, args ...string) (*ExecResult, error) { + if !s.created { + if err := s.Create(ctx); err != nil { + return nil, err + } + } + + start := time.Now() + + if s.config.Timeout > 0 { + var cancel context.CancelFunc + ctx, cancel = context.WithTimeout(ctx, s.config.Timeout) + defer cancel() + } + + dockerArgs := []string{"exec"} + if stdin != nil { + dockerArgs = append(dockerArgs, "-i") + } + dockerArgs = append(dockerArgs, s.containerName(), command) + dockerArgs = append(dockerArgs, args...) + + cmd := exec.CommandContext(ctx, "docker", dockerArgs...) + + var stdout, stderr bytes.Buffer + cmd.Stdout = &stdout + cmd.Stderr = &stderr + + if stdin != nil { + cmd.Stdin = stdin + } + + err := cmd.Run() + duration := time.Since(start) + + result := &ExecResult{ + Stdout: stdout.String(), + Stderr: stderr.String(), + Duration: duration, + } + + if ctx.Err() == context.DeadlineExceeded { + result.Killed = true + result.KillReason = "timeout" + result.ExitCode = -1 + return result, nil + } + + if err != nil { + if exitErr, ok := err.(*exec.ExitError); ok { + result.ExitCode = exitErr.ExitCode() + } else { + return nil, err + } + } + + return result, nil +} + +func (s *DockerSandbox) ReadFile(ctx context.Context, path string) ([]byte, error) { + if !s.created { + if err := s.Create(ctx); err != nil { + return nil, err + } + } + + containerPath := s.resolveContainerPath(path) + + cmd := exec.CommandContext(ctx, "docker", "exec", s.containerName(), "cat", containerPath) + output, err := cmd.Output() + if err != nil { + return nil, fmt.Errorf("failed to read file %s: %w", containerPath, err) + } + + return output, nil +} + +func (s *DockerSandbox) WriteFile(ctx context.Context, path string, content []byte, mode uint32) error { + if !s.created { + if err := s.Create(ctx); err != nil { + return err + } + } + + containerPath := s.resolveContainerPath(path) + + // Use docker cp approach which works with --read-only containers + // Write content to temp file on host, then docker cp into container + tmpFile, err := os.CreateTemp("", "docker-sandbox-write-*") + if err != nil { + return fmt.Errorf("failed to create temp file: %w", err) + } + tmpPath := tmpFile.Name() + defer os.Remove(tmpPath) + + if _, err := tmpFile.Write(content); err != nil { + tmpFile.Close() + return fmt.Errorf("failed to write temp file: %w", err) + } + tmpFile.Close() + + // Set permissions on host file before copying (since chmod in container may fail with --cap-drop ALL) + if mode != 0 { + if err := os.Chmod(tmpPath, os.FileMode(mode)); err != nil { + return fmt.Errorf("failed to set file permissions: %w", err) + } + } + + // Ensure the directory exists in the container (workspace is writable via bind mount) + dir := filepath.Dir(containerPath) + if dir != "/" && dir != "." { + mkdirCmd := exec.CommandContext(ctx, "docker", "exec", s.containerName(), "mkdir", "-p", dir) + // Ignore error - directory might already exist or be read-only (but workspace is writable) + mkdirCmd.Run() + } + + // Use docker cp to copy file into container (preserves permissions from host) + cpCmd := exec.CommandContext(ctx, "docker", "cp", tmpPath, fmt.Sprintf("%s:%s", s.containerName(), containerPath)) + if output, err := cpCmd.CombinedOutput(); err != nil { + return fmt.Errorf("failed to write file %s: %w, output: %s", containerPath, err, string(output)) + } + + return nil +} + +func (s *DockerSandbox) DeleteFile(ctx context.Context, path string) error { + if !s.created { + return nil + } + + containerPath := s.resolveContainerPath(path) + + // If workspace is a bind mount, we can delete from host side + // (more reliable with --cap-drop ALL and --read-only) + if s.config.WorkspacePath != "" && strings.HasPrefix(containerPath, "/workspace") { + relativePath := strings.TrimPrefix(containerPath, "/workspace/") + hostPath := filepath.Join(s.config.WorkspacePath, relativePath) + return os.Remove(hostPath) + } + + // Fallback to docker exec rm for non-workspace paths + cmd := exec.CommandContext(ctx, "docker", "exec", s.containerName(), "rm", "-f", containerPath) + return cmd.Run() +} + +func (s *DockerSandbox) ListFiles(ctx context.Context, path string) ([]FileInfo, error) { + if !s.created { + if err := s.Create(ctx); err != nil { + return nil, err + } + } + + containerPath := s.resolveContainerPath(path) + + cmd := exec.CommandContext(ctx, "docker", "exec", s.containerName(), + "find", containerPath, "-maxdepth", "1", "-printf", "%f\\t%s\\t%m\\t%T@\\t%y\\n") + + output, err := cmd.Output() + if err != nil { + return nil, fmt.Errorf("failed to list files in %s: %w", containerPath, err) + } + + var files []FileInfo + for _, line := range strings.Split(strings.TrimSpace(string(output)), "\n") { + if line == "" { + continue + } + parts := strings.Split(line, "\t") + if len(parts) < 5 { + continue + } + + name := parts[0] + if name == filepath.Base(containerPath) { + continue + } + + var size int64 + fmt.Sscanf(parts[1], "%d", &size) + + var mode uint32 + fmt.Sscanf(parts[2], "%o", &mode) + + var modTimeUnix float64 + fmt.Sscanf(parts[3], "%f", &modTimeUnix) + modTime := time.Unix(int64(modTimeUnix), 0) + + isDir := parts[4] == "d" + + files = append(files, FileInfo{ + Name: name, + Size: size, + Mode: mode, + ModTime: modTime, + IsDir: isDir, + }) + } + + return files, nil +} + +func (s *DockerSandbox) FileExists(ctx context.Context, path string) (bool, error) { + if !s.created { + if err := s.Create(ctx); err != nil { + return false, err + } + } + + containerPath := s.resolveContainerPath(path) + cmd := exec.CommandContext(ctx, "docker", "exec", s.containerName(), "test", "-e", containerPath) + err := cmd.Run() + if err == nil { + return true, nil + } + if _, ok := err.(*exec.ExitError); ok { + return false, nil + } + return false, err +} + +func (s *DockerSandbox) CopyIn(ctx context.Context, hostPath, sandboxPath string) error { + if !s.created { + if err := s.Create(ctx); err != nil { + return err + } + } + + containerPath := s.resolveContainerPath(sandboxPath) + cmd := exec.CommandContext(ctx, "docker", "cp", hostPath, fmt.Sprintf("%s:%s", s.containerName(), containerPath)) + if output, err := cmd.CombinedOutput(); err != nil { + return fmt.Errorf("failed to copy %s to container: %w, output: %s", hostPath, err, string(output)) + } + return nil +} + +func (s *DockerSandbox) CopyOut(ctx context.Context, sandboxPath, hostPath string) error { + if !s.created { + if err := s.Create(ctx); err != nil { + return err + } + } + + containerPath := s.resolveContainerPath(sandboxPath) + + if err := os.MkdirAll(filepath.Dir(hostPath), 0755); err != nil { + return err + } + + cmd := exec.CommandContext(ctx, "docker", "cp", fmt.Sprintf("%s:%s", s.containerName(), containerPath), hostPath) + if output, err := cmd.CombinedOutput(); err != nil { + return fmt.Errorf("failed to copy %s from container: %w, output: %s", containerPath, err, string(output)) + } + return nil +} + +func (s *DockerSandbox) GetMetrics(ctx context.Context) (*Metrics, error) { + if !s.created { + return &Metrics{}, nil + } + + cmd := exec.CommandContext(ctx, "docker", "stats", "--no-stream", "--format", + "{{.CPUPerc}}\t{{.MemUsage}}\t{{.NetIO}}\t{{.BlockIO}}\t{{.PIDs}}", s.containerName()) + + output, err := cmd.Output() + if err != nil { + return nil, fmt.Errorf("failed to get container stats: %w", err) + } + + metrics := &Metrics{} + parts := strings.Split(strings.TrimSpace(string(output)), "\t") + if len(parts) >= 1 { + cpuStr := strings.TrimSuffix(parts[0], "%") + fmt.Sscanf(cpuStr, "%f", &metrics.CPUUsage) + } + if len(parts) >= 5 { + var pids int + fmt.Sscanf(parts[4], "%d", &pids) + metrics.ProcessCount = pids + } + + return metrics, nil +} + +func (s *DockerSandbox) Destroy(ctx context.Context) error { + if !s.created { + return nil + } + + cmd := exec.CommandContext(ctx, "docker", "rm", "-f", s.containerName()) + if err := cmd.Run(); err != nil { + return fmt.Errorf("failed to remove container: %w", err) + } + + s.created = false + return nil +} + +func (s *DockerSandbox) ID() string { + return s.id +} + +func (s *DockerSandbox) Config() *Config { + return &s.config +} + +func (s *DockerSandbox) containerName() string { + return fmt.Sprintf("station-sandbox-%s", s.id[:8]) +} + +func (s *DockerSandbox) resolveContainerPath(path string) string { + if filepath.IsAbs(path) { + return path + } + return filepath.Join("/workspace", path) +} + +func (s *DockerSandbox) pullImage(ctx context.Context) error { + auth := s.config.RegistryAuth + + if auth == nil { + cmd := exec.CommandContext(ctx, "docker", "pull", s.config.Image) + if output, err := cmd.CombinedOutput(); err != nil { + return fmt.Errorf("docker pull failed: %w, output: %s", err, string(output)) + } + return nil + } + + if auth.DockerConfigPath != "" { + cmd := exec.CommandContext(ctx, "docker", "--config", filepath.Dir(auth.DockerConfigPath), "pull", s.config.Image) + if output, err := cmd.CombinedOutput(); err != nil { + return fmt.Errorf("docker pull with config failed: %w, output: %s", err, string(output)) + } + return nil + } + + if auth.Username != "" && auth.Password != "" { + loginCmd := exec.CommandContext(ctx, "docker", "login", + "-u", auth.Username, + "--password-stdin", + auth.ServerAddress) + loginCmd.Stdin = strings.NewReader(auth.Password) + if output, err := loginCmd.CombinedOutput(); err != nil { + return fmt.Errorf("docker login failed: %w, output: %s", err, string(output)) + } + + pullCmd := exec.CommandContext(ctx, "docker", "pull", s.config.Image) + if output, err := pullCmd.CombinedOutput(); err != nil { + return fmt.Errorf("docker pull failed: %w, output: %s", err, string(output)) + } + return nil + } + + if auth.IdentityToken != "" { + authConfig := map[string]string{ + "identitytoken": auth.IdentityToken, + } + if auth.ServerAddress != "" { + authConfig["serveraddress"] = auth.ServerAddress + } + authJSON, _ := json.Marshal(authConfig) + authB64 := base64.URLEncoding.EncodeToString(authJSON) + + cmd := exec.CommandContext(ctx, "docker", "pull", s.config.Image) + cmd.Env = append(os.Environ(), "DOCKER_AUTH_CONFIG="+authB64) + if output, err := cmd.CombinedOutput(); err != nil { + return fmt.Errorf("docker pull with identity token failed: %w, output: %s", err, string(output)) + } + return nil + } + + cmd := exec.CommandContext(ctx, "docker", "pull", s.config.Image) + if output, err := cmd.CombinedOutput(); err != nil { + return fmt.Errorf("docker pull failed: %w, output: %s", err, string(output)) + } + return nil +} diff --git a/pkg/harness/sandbox/docker_test.go b/pkg/harness/sandbox/docker_test.go new file mode 100644 index 00000000..ad60c75d --- /dev/null +++ b/pkg/harness/sandbox/docker_test.go @@ -0,0 +1,276 @@ +package sandbox + +import ( + "context" + "os" + "os/exec" + "path/filepath" + "testing" + "time" +) + +func dockerAvailable() bool { + cmd := exec.Command("docker", "info") + return cmd.Run() == nil +} + +func TestDockerSandbox_Create(t *testing.T) { + if !dockerAvailable() { + t.Skip("Docker not available") + } + + tmpDir := t.TempDir() + cfg := Config{ + Mode: ModeDocker, + Image: "alpine:latest", + WorkspacePath: filepath.Join(tmpDir, "workspace"), + Timeout: 5 * time.Minute, + } + + sb, err := NewDockerSandbox(cfg) + if err != nil { + t.Fatalf("NewDockerSandbox failed: %v", err) + } + + ctx := context.Background() + if err := sb.Create(ctx); err != nil { + t.Fatalf("Create failed: %v", err) + } + defer sb.Destroy(ctx) + + if !sb.created { + t.Error("sandbox should be marked as created") + } +} + +func TestDockerSandbox_Exec(t *testing.T) { + if !dockerAvailable() { + t.Skip("Docker not available") + } + + cfg := Config{ + Mode: ModeDocker, + Image: "alpine:latest", + Timeout: 5 * time.Minute, + } + + sb, _ := NewDockerSandbox(cfg) + ctx := context.Background() + defer sb.Destroy(ctx) + + result, err := sb.Exec(ctx, "echo", "hello from docker") + if err != nil { + t.Fatalf("Exec failed: %v", err) + } + + if result.ExitCode != 0 { + t.Errorf("expected exit code 0, got %d", result.ExitCode) + } + + if result.Stdout != "hello from docker\n" { + t.Errorf("expected stdout 'hello from docker\\n', got %q", result.Stdout) + } +} + +func TestDockerSandbox_CopyOut_AutoCreatesContainer(t *testing.T) { + if !dockerAvailable() { + t.Skip("Docker not available") + } + + tmpDir := t.TempDir() + cfg := Config{ + Mode: ModeDocker, + Image: "alpine:latest", + WorkspacePath: filepath.Join(tmpDir, "workspace"), + Timeout: 5 * time.Minute, + } + + sb, _ := NewDockerSandbox(cfg) + ctx := context.Background() + defer sb.Destroy(ctx) + + if err := sb.WriteFile(ctx, "test.txt", []byte("test content"), 0644); err != nil { + t.Fatalf("WriteFile failed: %v", err) + } + + destFile := filepath.Join(tmpDir, "dest.txt") + if err := sb.CopyOut(ctx, "test.txt", destFile); err != nil { + t.Fatalf("CopyOut failed: %v", err) + } + + content, err := os.ReadFile(destFile) + if err != nil { + t.Fatalf("failed to read dest file: %v", err) + } + if string(content) != "test content" { + t.Errorf("content mismatch: got %q", content) + } +} + +func TestDockerSandbox_RegistryAuth_Config(t *testing.T) { + cfg := Config{ + Mode: ModeDocker, + Image: "private-registry.example.com/myimage:latest", + RegistryAuth: &RegistryAuthConfig{ + Username: "testuser", + Password: "testpass", + ServerAddress: "private-registry.example.com", + }, + } + + sb, err := NewDockerSandbox(cfg) + if err != nil { + t.Fatalf("NewDockerSandbox failed: %v", err) + } + + if sb.config.RegistryAuth == nil { + t.Fatal("RegistryAuth should be set") + } + + if sb.config.RegistryAuth.Username != "testuser" { + t.Errorf("expected username 'testuser', got %q", sb.config.RegistryAuth.Username) + } + + if sb.config.RegistryAuth.ServerAddress != "private-registry.example.com" { + t.Errorf("expected server 'private-registry.example.com', got %q", sb.config.RegistryAuth.ServerAddress) + } +} + +func TestDockerSandbox_RegistryAuth_DockerConfig(t *testing.T) { + tmpDir := t.TempDir() + configPath := filepath.Join(tmpDir, "config.json") + configContent := `{"auths":{"https://index.docker.io/v1/":{"auth":"dGVzdDp0ZXN0"}}}` + os.WriteFile(configPath, []byte(configContent), 0600) + + cfg := Config{ + Mode: ModeDocker, + Image: "alpine:latest", + RegistryAuth: &RegistryAuthConfig{ + DockerConfigPath: configPath, + }, + } + + sb, err := NewDockerSandbox(cfg) + if err != nil { + t.Fatalf("NewDockerSandbox failed: %v", err) + } + + if sb.config.RegistryAuth.DockerConfigPath != configPath { + t.Errorf("expected DockerConfigPath %q, got %q", configPath, sb.config.RegistryAuth.DockerConfigPath) + } +} + +func TestDockerSandbox_RegistryAuth_IdentityToken(t *testing.T) { + cfg := Config{ + Mode: ModeDocker, + Image: "123456789.dkr.ecr.us-east-1.amazonaws.com/myapp:latest", + RegistryAuth: &RegistryAuthConfig{ + IdentityToken: "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9...", + ServerAddress: "123456789.dkr.ecr.us-east-1.amazonaws.com", + }, + } + + sb, err := NewDockerSandbox(cfg) + if err != nil { + t.Fatalf("NewDockerSandbox failed: %v", err) + } + + if sb.config.RegistryAuth.IdentityToken == "" { + t.Error("IdentityToken should be set") + } +} + +func TestDockerSandbox_pullImage_NoAuth(t *testing.T) { + if !dockerAvailable() { + t.Skip("Docker not available") + } + + cfg := Config{ + Mode: ModeDocker, + Image: "alpine:latest", + Timeout: 5 * time.Minute, + } + + sb, _ := NewDockerSandbox(cfg) + ctx := context.Background() + + err := sb.pullImage(ctx) + if err != nil { + t.Fatalf("pullImage failed: %v", err) + } +} + +func TestDockerSandbox_Destroy(t *testing.T) { + if !dockerAvailable() { + t.Skip("Docker not available") + } + + cfg := Config{ + Mode: ModeDocker, + Image: "alpine:latest", + Timeout: 5 * time.Minute, + } + + sb, _ := NewDockerSandbox(cfg) + ctx := context.Background() + + if err := sb.Create(ctx); err != nil { + t.Fatalf("Create failed: %v", err) + } + + if err := sb.Destroy(ctx); err != nil { + t.Fatalf("Destroy failed: %v", err) + } + + if sb.created { + t.Error("sandbox should be marked as not created after destroy") + } +} + +func TestDockerSandbox_FileOperations(t *testing.T) { + if !dockerAvailable() { + t.Skip("Docker not available") + } + + tmpDir := t.TempDir() + cfg := Config{ + Mode: ModeDocker, + Image: "alpine:latest", + WorkspacePath: filepath.Join(tmpDir, "workspace"), + Timeout: 5 * time.Minute, + } + + sb, _ := NewDockerSandbox(cfg) + ctx := context.Background() + defer sb.Destroy(ctx) + + content := []byte("docker file test") + if err := sb.WriteFile(ctx, "test.txt", content, 0644); err != nil { + t.Fatalf("WriteFile failed: %v", err) + } + + exists, err := sb.FileExists(ctx, "test.txt") + if err != nil { + t.Fatalf("FileExists failed: %v", err) + } + if !exists { + t.Error("file should exist after write") + } + + readContent, err := sb.ReadFile(ctx, "test.txt") + if err != nil { + t.Fatalf("ReadFile failed: %v", err) + } + if string(readContent) != string(content) { + t.Errorf("content mismatch: got %q, want %q", readContent, content) + } + + if err := sb.DeleteFile(ctx, "test.txt"); err != nil { + t.Fatalf("DeleteFile failed: %v", err) + } + + exists, _ = sb.FileExists(ctx, "test.txt") + if exists { + t.Error("file should not exist after delete") + } +} diff --git a/pkg/harness/sandbox/e2b.go b/pkg/harness/sandbox/e2b.go new file mode 100644 index 00000000..ad399c42 --- /dev/null +++ b/pkg/harness/sandbox/e2b.go @@ -0,0 +1,520 @@ +package sandbox + +import ( + "context" + "fmt" + "io" + "os" + "strings" + "time" + + "station/pkg/harness/sandbox/e2b" +) + +type E2BSandbox struct { + cfg Config + client *e2b.Client + streamingClient *e2b.StreamingClient + sandboxID string + created bool +} + +type E2BConfig struct { + APIKey string + TemplateID string + TimeoutSec int +} + +func NewE2BSandbox(cfg Config, e2bCfg E2BConfig) (*E2BSandbox, error) { + return newE2BSandbox(cfg, e2bCfg, "") +} + +func NewE2BSandboxWithAPIURL(cfg Config, e2bCfg E2BConfig, apiURL string) (*E2BSandbox, error) { + return newE2BSandbox(cfg, e2bCfg, apiURL) +} + +func newE2BSandbox(cfg Config, e2bCfg E2BConfig, apiURL string) (*E2BSandbox, error) { + if e2bCfg.APIKey == "" { + e2bCfg.APIKey = os.Getenv("E2B_API_KEY") + } + if e2bCfg.APIKey == "" { + return nil, fmt.Errorf("E2B_API_KEY is required") + } + + if e2bCfg.TemplateID == "" { + e2bCfg.TemplateID = "base" + } + + if e2bCfg.TimeoutSec == 0 { + e2bCfg.TimeoutSec = int(cfg.Timeout.Seconds()) + if e2bCfg.TimeoutSec == 0 { + e2bCfg.TimeoutSec = 300 + } + } + + clientCfg := e2b.ClientConfig{ + APIKey: e2bCfg.APIKey, + Timeout: 60 * time.Second, + } + if apiURL != "" { + clientCfg.APIURL = apiURL + } + + client := e2b.NewClient(clientCfg) + + return &E2BSandbox{ + cfg: cfg, + client: client, + }, nil +} + +func (s *E2BSandbox) Create(ctx context.Context) error { + if s.created { + return nil + } + + templateID := "base" + if s.cfg.Image != "" { + templateID = s.cfg.Image + } + + timeoutSec := int(s.cfg.Timeout.Seconds()) + if timeoutSec == 0 { + timeoutSec = 300 + } + + info, err := s.client.CreateSandbox(ctx, e2b.CreateSandboxRequest{ + TemplateID: templateID, + Timeout: timeoutSec, + EnvVars: s.cfg.Environment, + }) + if err != nil { + return fmt.Errorf("create e2b sandbox: %w", err) + } + + s.sandboxID = info.SandboxID + s.created = true + + s.streamingClient = e2b.NewStreamingClient(s.client.EnvdURL(), info.EnvdAccessToken) + + if s.cfg.WorkspacePath != "" { + _ = s.client.MakeDir(ctx, s.cfg.WorkspacePath) + } + + return nil +} + +func (s *E2BSandbox) Exec(ctx context.Context, command string, args ...string) (*ExecResult, error) { + if !s.created { + if err := s.Create(ctx); err != nil { + return nil, err + } + } + + start := time.Now() + + // Build full command with working directory if specified + fullCmd := command + if len(args) > 0 { + for _, arg := range args { + fullCmd += " " + arg + } + } + + // Wrap with cd if workspace path is set + if s.cfg.WorkspacePath != "" { + fullCmd = fmt.Sprintf("cd %s && %s", s.cfg.WorkspacePath, fullCmd) + } + + result, err := s.client.Exec(ctx, "sh", "-c", fullCmd) + if err != nil { + return nil, fmt.Errorf("exec in e2b sandbox: %w", err) + } + + return &ExecResult{ + ExitCode: result.ExitCode, + Stdout: result.Stdout, + Stderr: result.Stderr, + Duration: time.Since(start), + }, nil +} + +func (s *E2BSandbox) ExecWithStdin(ctx context.Context, stdin io.Reader, command string, args ...string) (*ExecResult, error) { + // E2B doesn't have native stdin support via REST, so we write to a temp file and redirect + if !s.created { + if err := s.Create(ctx); err != nil { + return nil, err + } + } + + // Read stdin content + stdinContent, err := io.ReadAll(stdin) + if err != nil { + return nil, fmt.Errorf("read stdin: %w", err) + } + + // Write to temp file + tmpFile := "/tmp/.stdin_" + fmt.Sprintf("%d", time.Now().UnixNano()) + if err := s.client.WriteFile(ctx, tmpFile, stdinContent); err != nil { + return nil, fmt.Errorf("write stdin to temp file: %w", err) + } + defer s.client.DeleteFile(ctx, tmpFile) + + // Build command with stdin redirect + fullCmd := command + if len(args) > 0 { + for _, arg := range args { + fullCmd += " " + arg + } + } + fullCmd = fmt.Sprintf("cat %s | %s", tmpFile, fullCmd) + + return s.Exec(ctx, "sh", "-c", fullCmd) +} + +func (s *E2BSandbox) ReadFile(ctx context.Context, path string) ([]byte, error) { + if !s.created { + if err := s.Create(ctx); err != nil { + return nil, err + } + } + + return s.client.ReadFile(ctx, path) +} + +func (s *E2BSandbox) WriteFile(ctx context.Context, path string, content []byte, mode uint32) error { + if !s.created { + if err := s.Create(ctx); err != nil { + return err + } + } + + if err := s.client.WriteFile(ctx, path, content); err != nil { + return err + } + + // Set file mode if not default + if mode != 0 && mode != 0644 { + _, err := s.client.Exec(ctx, "chmod", fmt.Sprintf("%o", mode), path) + if err != nil { + // Non-fatal, log and continue + } + } + + return nil +} + +func (s *E2BSandbox) DeleteFile(ctx context.Context, path string) error { + if !s.created { + return nil // Nothing to delete if sandbox doesn't exist + } + + return s.client.DeleteFile(ctx, path) +} + +func (s *E2BSandbox) ListFiles(ctx context.Context, path string) ([]FileInfo, error) { + if !s.created { + if err := s.Create(ctx); err != nil { + return nil, err + } + } + + e2bFiles, err := s.client.ListDir(ctx, path) + if err != nil { + return nil, err + } + + files := make([]FileInfo, len(e2bFiles)) + for i, f := range e2bFiles { + var mode uint32 = 0644 + if f.IsDir { + mode = 0755 + } + + files[i] = FileInfo{ + Name: f.Name, + Size: f.Size, + Mode: mode, + ModTime: f.ModTime, + IsDir: f.IsDir, + } + } + + return files, nil +} + +func (s *E2BSandbox) FileExists(ctx context.Context, path string) (bool, error) { + if !s.created { + if err := s.Create(ctx); err != nil { + return false, err + } + } + + return s.client.FileExists(ctx, path) +} + +func (s *E2BSandbox) CopyIn(ctx context.Context, hostPath, sandboxPath string) error { + if !s.created { + if err := s.Create(ctx); err != nil { + return err + } + } + + // Read from host + content, err := os.ReadFile(hostPath) + if err != nil { + return fmt.Errorf("read host file: %w", err) + } + + // Write to sandbox + return s.client.WriteFile(ctx, sandboxPath, content) +} + +func (s *E2BSandbox) CopyOut(ctx context.Context, sandboxPath, hostPath string) error { + if !s.created { + if err := s.Create(ctx); err != nil { + return err + } + } + + // Read from sandbox + content, err := s.client.ReadFile(ctx, sandboxPath) + if err != nil { + return fmt.Errorf("read sandbox file: %w", err) + } + + // Write to host + return os.WriteFile(hostPath, content, 0644) +} + +func (s *E2BSandbox) GetMetrics(ctx context.Context) (*Metrics, error) { + if !s.created { + return &Metrics{}, nil + } + + // E2B doesn't expose detailed metrics, return basic info from ps + result, err := s.client.Exec(ctx, "ps", "aux", "--no-headers", "|", "wc", "-l") + if err != nil { + return &Metrics{}, nil + } + + var processCount int + fmt.Sscanf(result.Stdout, "%d", &processCount) + + return &Metrics{ + ProcessCount: processCount, + }, nil +} + +func (s *E2BSandbox) Destroy(ctx context.Context) error { + if !s.created { + return nil + } + + err := s.client.Kill(ctx) + s.created = false + s.sandboxID = "" + return err +} + +func (s *E2BSandbox) ID() string { + return s.sandboxID +} + +func (s *E2BSandbox) Config() *Config { + return &s.cfg +} + +func (s *E2BSandbox) ExecStream(ctx context.Context, opts ExecOptions, command string, args ...string) (ProcessHandle, error) { + if !s.created { + if err := s.Create(ctx); err != nil { + return nil, err + } + } + + cwd := opts.Cwd + if cwd == "" { + cwd = s.cfg.WorkspacePath + } + + env := make(map[string]string) + for k, v := range s.cfg.Environment { + env[k] = v + } + for k, v := range opts.Env { + env[k] = v + } + + if s.streamingClient != nil { + streamHandle, err := s.streamingClient.ExecStream(ctx, e2b.StreamExecOptions{ + Cmd: command, + Args: args, + Cwd: cwd, + Env: env, + OnStdout: opts.OnStdout, + OnStderr: opts.OnStderr, + }) + if err != nil { + return nil, fmt.Errorf("start streaming exec: %w", err) + } + + return &e2bStreamingHandle{ + handle: streamHandle, + start: time.Now(), + onStdout: opts.OnStdout, + onStderr: opts.OnStderr, + }, nil + } + + fullCmd := command + for _, arg := range args { + fullCmd += " " + arg + } + if cwd != "" { + fullCmd = fmt.Sprintf("cd %s && %s", cwd, fullCmd) + } + + handle := &e2bProcessHandle{ + sandbox: s, + ctx: ctx, + command: fullCmd, + onStdout: opts.OnStdout, + onStderr: opts.OnStderr, + done: make(chan struct{}), + start: time.Now(), + } + + go handle.run() + + return handle, nil +} + +func (s *E2BSandbox) ListProcesses(ctx context.Context) ([]ProcessInfo, error) { + result, err := s.client.Exec(ctx, "ps", "-eo", "pid,cmd", "--no-headers") + if err != nil { + return nil, err + } + + var processes []ProcessInfo + lines := strings.Split(result.Stdout, "\n") + for _, line := range lines { + line = strings.TrimSpace(line) + if line == "" { + continue + } + fields := strings.Fields(line) + if len(fields) < 2 { + continue + } + var pid int + fmt.Sscanf(fields[0], "%d", &pid) + processes = append(processes, ProcessInfo{ + PID: pid, + Command: strings.Join(fields[1:], " "), + }) + } + + return processes, nil +} + +func (s *E2BSandbox) KillProcess(ctx context.Context, pid int) error { + _, err := s.client.Exec(ctx, "kill", "-9", fmt.Sprintf("%d", pid)) + return err +} + +var _ Sandbox = (*E2BSandbox)(nil) +var _ StreamingSandbox = (*E2BSandbox)(nil) + +type e2bProcessHandle struct { + sandbox *E2BSandbox + ctx context.Context + command string + onStdout func([]byte) + onStderr func([]byte) + done chan struct{} + start time.Time + result *ExecResult + err error +} + +func (h *e2bProcessHandle) PID() int { + return 0 +} + +func (h *e2bProcessHandle) Wait() (*ExecResult, error) { + <-h.done + return h.result, h.err +} + +func (h *e2bProcessHandle) Kill() error { + return fmt.Errorf("kill not supported for E2B streaming exec") +} + +func (h *e2bProcessHandle) SendStdin(data []byte) error { + return fmt.Errorf("stdin not supported for E2B streaming exec") +} + +func (h *e2bProcessHandle) run() { + defer close(h.done) + + execResult, err := h.sandbox.client.Exec(h.ctx, "sh", "-c", h.command) + if err != nil { + h.err = err + return + } + + if h.onStdout != nil && execResult.Stdout != "" { + h.onStdout([]byte(execResult.Stdout)) + } + if h.onStderr != nil && execResult.Stderr != "" { + h.onStderr([]byte(execResult.Stderr)) + } + + h.result = &ExecResult{ + ExitCode: execResult.ExitCode, + Stdout: execResult.Stdout, + Stderr: execResult.Stderr, + Duration: time.Since(h.start), + } +} + +type e2bStreamingHandle struct { + handle *e2b.StreamExecHandle + start time.Time + onStdout func([]byte) + onStderr func([]byte) + stdout []byte + stderr []byte +} + +func (h *e2bStreamingHandle) PID() int { + return int(h.handle.PID()) +} + +func (h *e2bStreamingHandle) Wait() (*ExecResult, error) { + exitCode, exitError, err := h.handle.Wait() + if err != nil { + return nil, err + } + + result := &ExecResult{ + ExitCode: int(exitCode), + Stdout: string(h.stdout), + Stderr: string(h.stderr), + Duration: time.Since(h.start), + } + + if exitError != "" { + result.Stderr += exitError + } + + return result, nil +} + +func (h *e2bStreamingHandle) Kill() error { + return h.handle.Close() +} + +func (h *e2bStreamingHandle) SendStdin(data []byte) error { + return fmt.Errorf("stdin via streaming handle not yet implemented") +} diff --git a/pkg/harness/sandbox/e2b/client.go b/pkg/harness/sandbox/e2b/client.go new file mode 100644 index 00000000..8255689f --- /dev/null +++ b/pkg/harness/sandbox/e2b/client.go @@ -0,0 +1,401 @@ +package e2b + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "mime/multipart" + "net/http" + "net/url" + "strings" + "time" +) + +const ( + DefaultAPIURL = "https://api.e2b.dev" + DefaultDomain = "e2b.dev" + EnvdPort = 49983 + DefaultTimeoutSec = 300 +) + +type Client struct { + apiKey string + apiURL string + domain string + httpClient *http.Client + sandboxID string + accessToken string + envdURL string + sandboxDomain string +} + +type ClientConfig struct { + APIKey string + APIURL string + Domain string + Timeout time.Duration +} + +func NewClient(cfg ClientConfig) *Client { + if cfg.APIURL == "" { + cfg.APIURL = DefaultAPIURL + } + if cfg.Domain == "" { + cfg.Domain = DefaultDomain + } + if cfg.Timeout == 0 { + cfg.Timeout = 60 * time.Second + } + + return &Client{ + apiKey: cfg.APIKey, + apiURL: cfg.APIURL, + domain: cfg.Domain, + httpClient: &http.Client{ + Timeout: cfg.Timeout, + }, + } +} + +type CreateSandboxRequest struct { + TemplateID string `json:"templateID"` + Timeout int `json:"timeout"` + Metadata map[string]string `json:"metadata,omitempty"` + EnvVars map[string]string `json:"envVars,omitempty"` +} + +type SandboxInfo struct { + SandboxID string `json:"sandboxID"` + TemplateID string `json:"templateID"` + Domain string `json:"domain"` + EnvdAccessToken string `json:"envdAccessToken"` + TrafficAccessToken string `json:"trafficAccessToken,omitempty"` + EnvdVersion string `json:"envdVersion"` + StartedAt string `json:"startedAt"` + EndAt string `json:"endAt"` +} + +func (c *Client) CreateSandbox(ctx context.Context, req CreateSandboxRequest) (*SandboxInfo, error) { + if req.Timeout == 0 { + req.Timeout = DefaultTimeoutSec + } + + body, err := json.Marshal(req) + if err != nil { + return nil, fmt.Errorf("marshal request: %w", err) + } + + httpReq, err := http.NewRequestWithContext(ctx, "POST", c.apiURL+"/sandboxes", bytes.NewReader(body)) + if err != nil { + return nil, fmt.Errorf("create request: %w", err) + } + + httpReq.Header.Set("Content-Type", "application/json") + httpReq.Header.Set("X-API-KEY", c.apiKey) + + resp, err := c.httpClient.Do(httpReq) + if err != nil { + return nil, fmt.Errorf("execute request: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusCreated && resp.StatusCode != http.StatusOK { + respBody, _ := io.ReadAll(resp.Body) + return nil, fmt.Errorf("create sandbox failed: status %d, body: %s", resp.StatusCode, string(respBody)) + } + + var info SandboxInfo + if err := json.NewDecoder(resp.Body).Decode(&info); err != nil { + return nil, fmt.Errorf("decode response: %w", err) + } + + c.sandboxID = info.SandboxID + c.accessToken = info.EnvdAccessToken + c.sandboxDomain = info.Domain + if c.sandboxDomain == "" { + c.sandboxDomain = c.domain + } + c.envdURL = fmt.Sprintf("https://%d-%s.%s", EnvdPort, info.SandboxID, c.sandboxDomain) + + return &info, nil +} + +func (c *Client) Connect(sandboxID string, accessToken string, domain string) { + c.sandboxID = sandboxID + c.accessToken = accessToken + if domain == "" { + domain = c.domain + } + c.sandboxDomain = domain + c.envdURL = fmt.Sprintf("https://%d-%s.%s", EnvdPort, sandboxID, domain) +} + +func (c *Client) Kill(ctx context.Context) error { + if c.sandboxID == "" { + return nil + } + + httpReq, err := http.NewRequestWithContext(ctx, "DELETE", c.apiURL+"/sandboxes/"+c.sandboxID, nil) + if err != nil { + return fmt.Errorf("create request: %w", err) + } + + httpReq.Header.Set("X-API-KEY", c.apiKey) + + resp, err := c.httpClient.Do(httpReq) + if err != nil { + return fmt.Errorf("execute request: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusNoContent && resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusNotFound { + respBody, _ := io.ReadAll(resp.Body) + return fmt.Errorf("kill sandbox failed: status %d, body: %s", resp.StatusCode, string(respBody)) + } + + return nil +} + +type ExecResult struct { + ExitCode int + Stdout string + Stderr string + Error string +} + +func (c *Client) Exec(ctx context.Context, command string, args ...string) (*ExecResult, error) { + fullCmd := command + if len(args) > 0 { + fullCmd = command + " " + strings.Join(args, " ") + } + + reqBody := map[string]interface{}{ + "cmd": fullCmd, + "wait": true, + } + + body, err := json.Marshal(reqBody) + if err != nil { + return nil, fmt.Errorf("marshal request: %w", err) + } + + httpReq, err := http.NewRequestWithContext(ctx, "POST", c.envdURL+"/commands", bytes.NewReader(body)) + if err != nil { + return nil, fmt.Errorf("create request: %w", err) + } + + httpReq.Header.Set("Content-Type", "application/json") + httpReq.Header.Set("X-Access-Token", c.accessToken) + + resp, err := c.httpClient.Do(httpReq) + if err != nil { + return nil, fmt.Errorf("execute request: %w", err) + } + defer resp.Body.Close() + + respBody, err := io.ReadAll(resp.Body) + if err != nil { + return nil, fmt.Errorf("read response: %w", err) + } + + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("exec failed: status %d, body: %s", resp.StatusCode, string(respBody)) + } + + var result struct { + ExitCode int `json:"exitCode"` + Stdout string `json:"stdout"` + Stderr string `json:"stderr"` + Error string `json:"error,omitempty"` + } + + if err := json.Unmarshal(respBody, &result); err != nil { + return nil, fmt.Errorf("decode response: %w", err) + } + + return &ExecResult{ + ExitCode: result.ExitCode, + Stdout: result.Stdout, + Stderr: result.Stderr, + Error: result.Error, + }, nil +} + +func (c *Client) ReadFile(ctx context.Context, path string) ([]byte, error) { + reqURL := c.envdURL + "/files?" + url.Values{"path": {path}}.Encode() + + req, err := http.NewRequestWithContext(ctx, "GET", reqURL, nil) + if err != nil { + return nil, fmt.Errorf("create request: %w", err) + } + + req.Header.Set("X-Access-Token", c.accessToken) + + resp, err := c.httpClient.Do(req) + if err != nil { + return nil, fmt.Errorf("execute request: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + body, _ := io.ReadAll(resp.Body) + return nil, fmt.Errorf("read file failed: status %d, body: %s", resp.StatusCode, string(body)) + } + + return io.ReadAll(resp.Body) +} + +func (c *Client) WriteFile(ctx context.Context, path string, content []byte) error { + var buf bytes.Buffer + writer := multipart.NewWriter(&buf) + + part, err := writer.CreateFormFile("file", path) + if err != nil { + return fmt.Errorf("create form file: %w", err) + } + + if _, err := part.Write(content); err != nil { + return fmt.Errorf("write content: %w", err) + } + + if err := writer.Close(); err != nil { + return fmt.Errorf("close writer: %w", err) + } + + reqURL := c.envdURL + "/files?" + url.Values{"path": {path}}.Encode() + + req, err := http.NewRequestWithContext(ctx, "POST", reqURL, &buf) + if err != nil { + return fmt.Errorf("create request: %w", err) + } + + req.Header.Set("Content-Type", writer.FormDataContentType()) + req.Header.Set("X-Access-Token", c.accessToken) + + resp, err := c.httpClient.Do(req) + if err != nil { + return fmt.Errorf("execute request: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusCreated { + body, _ := io.ReadAll(resp.Body) + return fmt.Errorf("write file failed: status %d, body: %s", resp.StatusCode, string(body)) + } + + return nil +} + +func (c *Client) DeleteFile(ctx context.Context, path string) error { + reqURL := c.envdURL + "/files?" + url.Values{"path": {path}}.Encode() + + req, err := http.NewRequestWithContext(ctx, "DELETE", reqURL, nil) + if err != nil { + return fmt.Errorf("create request: %w", err) + } + + req.Header.Set("X-Access-Token", c.accessToken) + + resp, err := c.httpClient.Do(req) + if err != nil { + return fmt.Errorf("execute request: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusNoContent && resp.StatusCode != http.StatusNotFound { + body, _ := io.ReadAll(resp.Body) + return fmt.Errorf("delete file failed: status %d, body: %s", resp.StatusCode, string(body)) + } + + return nil +} + +func (c *Client) MakeDir(ctx context.Context, path string) error { + return c.WriteFile(ctx, path+"/.keep", []byte{}) +} + +type FileInfo struct { + Name string + Path string + Size int64 + IsDir bool + ModTime time.Time +} + +func (c *Client) ListDir(ctx context.Context, path string) ([]FileInfo, error) { + result, err := c.Exec(ctx, "ls", "-la", "--time-style=+%s", path) + if err != nil { + return nil, err + } + + if result.ExitCode != 0 { + return nil, fmt.Errorf("ls failed: %s", result.Stderr) + } + + var files []FileInfo + lines := strings.Split(result.Stdout, "\n") + for _, line := range lines { + line = strings.TrimSpace(line) + if line == "" || strings.HasPrefix(line, "total") { + continue + } + + fields := strings.Fields(line) + if len(fields) < 7 { + continue + } + + name := fields[len(fields)-1] + if name == "." || name == ".." { + continue + } + + isDir := strings.HasPrefix(fields[0], "d") + + var size int64 + fmt.Sscanf(fields[4], "%d", &size) + + var modTime time.Time + var timestamp int64 + if n, _ := fmt.Sscanf(fields[5], "%d", ×tamp); n == 1 { + modTime = time.Unix(timestamp, 0) + } + + files = append(files, FileInfo{ + Name: name, + Path: path + "/" + name, + Size: size, + IsDir: isDir, + ModTime: modTime, + }) + } + + return files, nil +} + +func (c *Client) FileExists(ctx context.Context, path string) (bool, error) { + result, err := c.Exec(ctx, "test", "-e", path, "&&", "echo", "exists") + if err != nil { + return false, err + } + + return strings.Contains(result.Stdout, "exists"), nil +} + +func (c *Client) SandboxID() string { + return c.sandboxID +} + +func (c *Client) EnvdURL() string { + return c.envdURL +} + +func (c *Client) SetEnvdURL(url string) { + c.envdURL = url +} + +func (c *Client) SetAPIURL(url string) { + c.apiURL = url +} diff --git a/pkg/harness/sandbox/e2b/proto/buf.gen.yaml b/pkg/harness/sandbox/e2b/proto/buf.gen.yaml new file mode 100644 index 00000000..5a3bb6a8 --- /dev/null +++ b/pkg/harness/sandbox/e2b/proto/buf.gen.yaml @@ -0,0 +1,8 @@ +version: v2 +plugins: + - remote: buf.build/protocolbuffers/go + out: gen + opt: paths=source_relative + - remote: buf.build/connectrpc/go + out: gen + opt: paths=source_relative diff --git a/pkg/harness/sandbox/e2b/proto/buf.yaml b/pkg/harness/sandbox/e2b/proto/buf.yaml new file mode 100644 index 00000000..f74da98a --- /dev/null +++ b/pkg/harness/sandbox/e2b/proto/buf.yaml @@ -0,0 +1,9 @@ +version: v2 +modules: + - path: . +lint: + use: + - STANDARD +breaking: + use: + - FILE diff --git a/pkg/harness/sandbox/e2b/proto/gen/process/process.pb.go b/pkg/harness/sandbox/e2b/proto/gen/process/process.pb.go new file mode 100644 index 00000000..bbbfe710 --- /dev/null +++ b/pkg/harness/sandbox/e2b/proto/gen/process/process.pb.go @@ -0,0 +1,1877 @@ +// Code generated by protoc-gen-go. DO NOT EDIT. +// versions: +// protoc-gen-go v1.36.11 +// protoc (unknown) +// source: process.proto + +package process + +import ( + protoreflect "google.golang.org/protobuf/reflect/protoreflect" + protoimpl "google.golang.org/protobuf/runtime/protoimpl" + reflect "reflect" + sync "sync" + unsafe "unsafe" +) + +const ( + // Verify that this generated code is sufficiently up-to-date. + _ = protoimpl.EnforceVersion(20 - protoimpl.MinVersion) + // Verify that runtime/protoimpl is sufficiently up-to-date. + _ = protoimpl.EnforceVersion(protoimpl.MaxVersion - 20) +) + +type Signal int32 + +const ( + Signal_SIGNAL_UNSPECIFIED Signal = 0 + Signal_SIGNAL_SIGTERM Signal = 15 + Signal_SIGNAL_SIGKILL Signal = 9 +) + +// Enum value maps for Signal. +var ( + Signal_name = map[int32]string{ + 0: "SIGNAL_UNSPECIFIED", + 15: "SIGNAL_SIGTERM", + 9: "SIGNAL_SIGKILL", + } + Signal_value = map[string]int32{ + "SIGNAL_UNSPECIFIED": 0, + "SIGNAL_SIGTERM": 15, + "SIGNAL_SIGKILL": 9, + } +) + +func (x Signal) Enum() *Signal { + p := new(Signal) + *p = x + return p +} + +func (x Signal) String() string { + return protoimpl.X.EnumStringOf(x.Descriptor(), protoreflect.EnumNumber(x)) +} + +func (Signal) Descriptor() protoreflect.EnumDescriptor { + return file_process_proto_enumTypes[0].Descriptor() +} + +func (Signal) Type() protoreflect.EnumType { + return &file_process_proto_enumTypes[0] +} + +func (x Signal) Number() protoreflect.EnumNumber { + return protoreflect.EnumNumber(x) +} + +// Deprecated: Use Signal.Descriptor instead. +func (Signal) EnumDescriptor() ([]byte, []int) { + return file_process_proto_rawDescGZIP(), []int{0} +} + +type PTY struct { + state protoimpl.MessageState `protogen:"open.v1"` + Size *PTY_Size `protobuf:"bytes,1,opt,name=size,proto3" json:"size,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *PTY) Reset() { + *x = PTY{} + mi := &file_process_proto_msgTypes[0] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *PTY) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*PTY) ProtoMessage() {} + +func (x *PTY) ProtoReflect() protoreflect.Message { + mi := &file_process_proto_msgTypes[0] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use PTY.ProtoReflect.Descriptor instead. +func (*PTY) Descriptor() ([]byte, []int) { + return file_process_proto_rawDescGZIP(), []int{0} +} + +func (x *PTY) GetSize() *PTY_Size { + if x != nil { + return x.Size + } + return nil +} + +type ProcessConfig struct { + state protoimpl.MessageState `protogen:"open.v1"` + Cmd string `protobuf:"bytes,1,opt,name=cmd,proto3" json:"cmd,omitempty"` + Args []string `protobuf:"bytes,2,rep,name=args,proto3" json:"args,omitempty"` + Envs map[string]string `protobuf:"bytes,3,rep,name=envs,proto3" json:"envs,omitempty" protobuf_key:"bytes,1,opt,name=key" protobuf_val:"bytes,2,opt,name=value"` + Cwd *string `protobuf:"bytes,4,opt,name=cwd,proto3,oneof" json:"cwd,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *ProcessConfig) Reset() { + *x = ProcessConfig{} + mi := &file_process_proto_msgTypes[1] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *ProcessConfig) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ProcessConfig) ProtoMessage() {} + +func (x *ProcessConfig) ProtoReflect() protoreflect.Message { + mi := &file_process_proto_msgTypes[1] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ProcessConfig.ProtoReflect.Descriptor instead. +func (*ProcessConfig) Descriptor() ([]byte, []int) { + return file_process_proto_rawDescGZIP(), []int{1} +} + +func (x *ProcessConfig) GetCmd() string { + if x != nil { + return x.Cmd + } + return "" +} + +func (x *ProcessConfig) GetArgs() []string { + if x != nil { + return x.Args + } + return nil +} + +func (x *ProcessConfig) GetEnvs() map[string]string { + if x != nil { + return x.Envs + } + return nil +} + +func (x *ProcessConfig) GetCwd() string { + if x != nil && x.Cwd != nil { + return *x.Cwd + } + return "" +} + +type ListRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *ListRequest) Reset() { + *x = ListRequest{} + mi := &file_process_proto_msgTypes[2] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *ListRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ListRequest) ProtoMessage() {} + +func (x *ListRequest) ProtoReflect() protoreflect.Message { + mi := &file_process_proto_msgTypes[2] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ListRequest.ProtoReflect.Descriptor instead. +func (*ListRequest) Descriptor() ([]byte, []int) { + return file_process_proto_rawDescGZIP(), []int{2} +} + +type ProcessInfo struct { + state protoimpl.MessageState `protogen:"open.v1"` + Config *ProcessConfig `protobuf:"bytes,1,opt,name=config,proto3" json:"config,omitempty"` + Pid uint32 `protobuf:"varint,2,opt,name=pid,proto3" json:"pid,omitempty"` + Tag *string `protobuf:"bytes,3,opt,name=tag,proto3,oneof" json:"tag,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *ProcessInfo) Reset() { + *x = ProcessInfo{} + mi := &file_process_proto_msgTypes[3] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *ProcessInfo) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ProcessInfo) ProtoMessage() {} + +func (x *ProcessInfo) ProtoReflect() protoreflect.Message { + mi := &file_process_proto_msgTypes[3] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ProcessInfo.ProtoReflect.Descriptor instead. +func (*ProcessInfo) Descriptor() ([]byte, []int) { + return file_process_proto_rawDescGZIP(), []int{3} +} + +func (x *ProcessInfo) GetConfig() *ProcessConfig { + if x != nil { + return x.Config + } + return nil +} + +func (x *ProcessInfo) GetPid() uint32 { + if x != nil { + return x.Pid + } + return 0 +} + +func (x *ProcessInfo) GetTag() string { + if x != nil && x.Tag != nil { + return *x.Tag + } + return "" +} + +type ListResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + Processes []*ProcessInfo `protobuf:"bytes,1,rep,name=processes,proto3" json:"processes,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *ListResponse) Reset() { + *x = ListResponse{} + mi := &file_process_proto_msgTypes[4] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *ListResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ListResponse) ProtoMessage() {} + +func (x *ListResponse) ProtoReflect() protoreflect.Message { + mi := &file_process_proto_msgTypes[4] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ListResponse.ProtoReflect.Descriptor instead. +func (*ListResponse) Descriptor() ([]byte, []int) { + return file_process_proto_rawDescGZIP(), []int{4} +} + +func (x *ListResponse) GetProcesses() []*ProcessInfo { + if x != nil { + return x.Processes + } + return nil +} + +type StartRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + Process *ProcessConfig `protobuf:"bytes,1,opt,name=process,proto3" json:"process,omitempty"` + Pty *PTY `protobuf:"bytes,2,opt,name=pty,proto3,oneof" json:"pty,omitempty"` + Tag *string `protobuf:"bytes,3,opt,name=tag,proto3,oneof" json:"tag,omitempty"` + Stdin *bool `protobuf:"varint,4,opt,name=stdin,proto3,oneof" json:"stdin,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *StartRequest) Reset() { + *x = StartRequest{} + mi := &file_process_proto_msgTypes[5] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *StartRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*StartRequest) ProtoMessage() {} + +func (x *StartRequest) ProtoReflect() protoreflect.Message { + mi := &file_process_proto_msgTypes[5] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use StartRequest.ProtoReflect.Descriptor instead. +func (*StartRequest) Descriptor() ([]byte, []int) { + return file_process_proto_rawDescGZIP(), []int{5} +} + +func (x *StartRequest) GetProcess() *ProcessConfig { + if x != nil { + return x.Process + } + return nil +} + +func (x *StartRequest) GetPty() *PTY { + if x != nil { + return x.Pty + } + return nil +} + +func (x *StartRequest) GetTag() string { + if x != nil && x.Tag != nil { + return *x.Tag + } + return "" +} + +func (x *StartRequest) GetStdin() bool { + if x != nil && x.Stdin != nil { + return *x.Stdin + } + return false +} + +type UpdateRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + Process *ProcessSelector `protobuf:"bytes,1,opt,name=process,proto3" json:"process,omitempty"` + Pty *PTY `protobuf:"bytes,2,opt,name=pty,proto3,oneof" json:"pty,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *UpdateRequest) Reset() { + *x = UpdateRequest{} + mi := &file_process_proto_msgTypes[6] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *UpdateRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*UpdateRequest) ProtoMessage() {} + +func (x *UpdateRequest) ProtoReflect() protoreflect.Message { + mi := &file_process_proto_msgTypes[6] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use UpdateRequest.ProtoReflect.Descriptor instead. +func (*UpdateRequest) Descriptor() ([]byte, []int) { + return file_process_proto_rawDescGZIP(), []int{6} +} + +func (x *UpdateRequest) GetProcess() *ProcessSelector { + if x != nil { + return x.Process + } + return nil +} + +func (x *UpdateRequest) GetPty() *PTY { + if x != nil { + return x.Pty + } + return nil +} + +type UpdateResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *UpdateResponse) Reset() { + *x = UpdateResponse{} + mi := &file_process_proto_msgTypes[7] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *UpdateResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*UpdateResponse) ProtoMessage() {} + +func (x *UpdateResponse) ProtoReflect() protoreflect.Message { + mi := &file_process_proto_msgTypes[7] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use UpdateResponse.ProtoReflect.Descriptor instead. +func (*UpdateResponse) Descriptor() ([]byte, []int) { + return file_process_proto_rawDescGZIP(), []int{7} +} + +type ProcessEvent struct { + state protoimpl.MessageState `protogen:"open.v1"` + // Types that are valid to be assigned to Event: + // + // *ProcessEvent_Start + // *ProcessEvent_Data + // *ProcessEvent_End + // *ProcessEvent_Keepalive + Event isProcessEvent_Event `protobuf_oneof:"event"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *ProcessEvent) Reset() { + *x = ProcessEvent{} + mi := &file_process_proto_msgTypes[8] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *ProcessEvent) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ProcessEvent) ProtoMessage() {} + +func (x *ProcessEvent) ProtoReflect() protoreflect.Message { + mi := &file_process_proto_msgTypes[8] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ProcessEvent.ProtoReflect.Descriptor instead. +func (*ProcessEvent) Descriptor() ([]byte, []int) { + return file_process_proto_rawDescGZIP(), []int{8} +} + +func (x *ProcessEvent) GetEvent() isProcessEvent_Event { + if x != nil { + return x.Event + } + return nil +} + +func (x *ProcessEvent) GetStart() *ProcessEvent_StartEvent { + if x != nil { + if x, ok := x.Event.(*ProcessEvent_Start); ok { + return x.Start + } + } + return nil +} + +func (x *ProcessEvent) GetData() *ProcessEvent_DataEvent { + if x != nil { + if x, ok := x.Event.(*ProcessEvent_Data); ok { + return x.Data + } + } + return nil +} + +func (x *ProcessEvent) GetEnd() *ProcessEvent_EndEvent { + if x != nil { + if x, ok := x.Event.(*ProcessEvent_End); ok { + return x.End + } + } + return nil +} + +func (x *ProcessEvent) GetKeepalive() *ProcessEvent_KeepAlive { + if x != nil { + if x, ok := x.Event.(*ProcessEvent_Keepalive); ok { + return x.Keepalive + } + } + return nil +} + +type isProcessEvent_Event interface { + isProcessEvent_Event() +} + +type ProcessEvent_Start struct { + Start *ProcessEvent_StartEvent `protobuf:"bytes,1,opt,name=start,proto3,oneof"` +} + +type ProcessEvent_Data struct { + Data *ProcessEvent_DataEvent `protobuf:"bytes,2,opt,name=data,proto3,oneof"` +} + +type ProcessEvent_End struct { + End *ProcessEvent_EndEvent `protobuf:"bytes,3,opt,name=end,proto3,oneof"` +} + +type ProcessEvent_Keepalive struct { + Keepalive *ProcessEvent_KeepAlive `protobuf:"bytes,4,opt,name=keepalive,proto3,oneof"` +} + +func (*ProcessEvent_Start) isProcessEvent_Event() {} + +func (*ProcessEvent_Data) isProcessEvent_Event() {} + +func (*ProcessEvent_End) isProcessEvent_Event() {} + +func (*ProcessEvent_Keepalive) isProcessEvent_Event() {} + +type StartResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + Event *ProcessEvent `protobuf:"bytes,1,opt,name=event,proto3" json:"event,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *StartResponse) Reset() { + *x = StartResponse{} + mi := &file_process_proto_msgTypes[9] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *StartResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*StartResponse) ProtoMessage() {} + +func (x *StartResponse) ProtoReflect() protoreflect.Message { + mi := &file_process_proto_msgTypes[9] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use StartResponse.ProtoReflect.Descriptor instead. +func (*StartResponse) Descriptor() ([]byte, []int) { + return file_process_proto_rawDescGZIP(), []int{9} +} + +func (x *StartResponse) GetEvent() *ProcessEvent { + if x != nil { + return x.Event + } + return nil +} + +type ConnectResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + Event *ProcessEvent `protobuf:"bytes,1,opt,name=event,proto3" json:"event,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *ConnectResponse) Reset() { + *x = ConnectResponse{} + mi := &file_process_proto_msgTypes[10] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *ConnectResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ConnectResponse) ProtoMessage() {} + +func (x *ConnectResponse) ProtoReflect() protoreflect.Message { + mi := &file_process_proto_msgTypes[10] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ConnectResponse.ProtoReflect.Descriptor instead. +func (*ConnectResponse) Descriptor() ([]byte, []int) { + return file_process_proto_rawDescGZIP(), []int{10} +} + +func (x *ConnectResponse) GetEvent() *ProcessEvent { + if x != nil { + return x.Event + } + return nil +} + +type SendInputRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + Process *ProcessSelector `protobuf:"bytes,1,opt,name=process,proto3" json:"process,omitempty"` + Input *ProcessInput `protobuf:"bytes,2,opt,name=input,proto3" json:"input,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *SendInputRequest) Reset() { + *x = SendInputRequest{} + mi := &file_process_proto_msgTypes[11] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *SendInputRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*SendInputRequest) ProtoMessage() {} + +func (x *SendInputRequest) ProtoReflect() protoreflect.Message { + mi := &file_process_proto_msgTypes[11] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use SendInputRequest.ProtoReflect.Descriptor instead. +func (*SendInputRequest) Descriptor() ([]byte, []int) { + return file_process_proto_rawDescGZIP(), []int{11} +} + +func (x *SendInputRequest) GetProcess() *ProcessSelector { + if x != nil { + return x.Process + } + return nil +} + +func (x *SendInputRequest) GetInput() *ProcessInput { + if x != nil { + return x.Input + } + return nil +} + +type SendInputResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *SendInputResponse) Reset() { + *x = SendInputResponse{} + mi := &file_process_proto_msgTypes[12] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *SendInputResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*SendInputResponse) ProtoMessage() {} + +func (x *SendInputResponse) ProtoReflect() protoreflect.Message { + mi := &file_process_proto_msgTypes[12] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use SendInputResponse.ProtoReflect.Descriptor instead. +func (*SendInputResponse) Descriptor() ([]byte, []int) { + return file_process_proto_rawDescGZIP(), []int{12} +} + +type ProcessInput struct { + state protoimpl.MessageState `protogen:"open.v1"` + // Types that are valid to be assigned to Input: + // + // *ProcessInput_Stdin + // *ProcessInput_Pty + Input isProcessInput_Input `protobuf_oneof:"input"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *ProcessInput) Reset() { + *x = ProcessInput{} + mi := &file_process_proto_msgTypes[13] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *ProcessInput) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ProcessInput) ProtoMessage() {} + +func (x *ProcessInput) ProtoReflect() protoreflect.Message { + mi := &file_process_proto_msgTypes[13] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ProcessInput.ProtoReflect.Descriptor instead. +func (*ProcessInput) Descriptor() ([]byte, []int) { + return file_process_proto_rawDescGZIP(), []int{13} +} + +func (x *ProcessInput) GetInput() isProcessInput_Input { + if x != nil { + return x.Input + } + return nil +} + +func (x *ProcessInput) GetStdin() []byte { + if x != nil { + if x, ok := x.Input.(*ProcessInput_Stdin); ok { + return x.Stdin + } + } + return nil +} + +func (x *ProcessInput) GetPty() []byte { + if x != nil { + if x, ok := x.Input.(*ProcessInput_Pty); ok { + return x.Pty + } + } + return nil +} + +type isProcessInput_Input interface { + isProcessInput_Input() +} + +type ProcessInput_Stdin struct { + Stdin []byte `protobuf:"bytes,1,opt,name=stdin,proto3,oneof"` +} + +type ProcessInput_Pty struct { + Pty []byte `protobuf:"bytes,2,opt,name=pty,proto3,oneof"` +} + +func (*ProcessInput_Stdin) isProcessInput_Input() {} + +func (*ProcessInput_Pty) isProcessInput_Input() {} + +type StreamInputRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + // Types that are valid to be assigned to Event: + // + // *StreamInputRequest_Start + // *StreamInputRequest_Data + // *StreamInputRequest_Keepalive + Event isStreamInputRequest_Event `protobuf_oneof:"event"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *StreamInputRequest) Reset() { + *x = StreamInputRequest{} + mi := &file_process_proto_msgTypes[14] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *StreamInputRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*StreamInputRequest) ProtoMessage() {} + +func (x *StreamInputRequest) ProtoReflect() protoreflect.Message { + mi := &file_process_proto_msgTypes[14] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use StreamInputRequest.ProtoReflect.Descriptor instead. +func (*StreamInputRequest) Descriptor() ([]byte, []int) { + return file_process_proto_rawDescGZIP(), []int{14} +} + +func (x *StreamInputRequest) GetEvent() isStreamInputRequest_Event { + if x != nil { + return x.Event + } + return nil +} + +func (x *StreamInputRequest) GetStart() *StreamInputRequest_StartEvent { + if x != nil { + if x, ok := x.Event.(*StreamInputRequest_Start); ok { + return x.Start + } + } + return nil +} + +func (x *StreamInputRequest) GetData() *StreamInputRequest_DataEvent { + if x != nil { + if x, ok := x.Event.(*StreamInputRequest_Data); ok { + return x.Data + } + } + return nil +} + +func (x *StreamInputRequest) GetKeepalive() *StreamInputRequest_KeepAlive { + if x != nil { + if x, ok := x.Event.(*StreamInputRequest_Keepalive); ok { + return x.Keepalive + } + } + return nil +} + +type isStreamInputRequest_Event interface { + isStreamInputRequest_Event() +} + +type StreamInputRequest_Start struct { + Start *StreamInputRequest_StartEvent `protobuf:"bytes,1,opt,name=start,proto3,oneof"` +} + +type StreamInputRequest_Data struct { + Data *StreamInputRequest_DataEvent `protobuf:"bytes,2,opt,name=data,proto3,oneof"` +} + +type StreamInputRequest_Keepalive struct { + Keepalive *StreamInputRequest_KeepAlive `protobuf:"bytes,3,opt,name=keepalive,proto3,oneof"` +} + +func (*StreamInputRequest_Start) isStreamInputRequest_Event() {} + +func (*StreamInputRequest_Data) isStreamInputRequest_Event() {} + +func (*StreamInputRequest_Keepalive) isStreamInputRequest_Event() {} + +type StreamInputResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *StreamInputResponse) Reset() { + *x = StreamInputResponse{} + mi := &file_process_proto_msgTypes[15] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *StreamInputResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*StreamInputResponse) ProtoMessage() {} + +func (x *StreamInputResponse) ProtoReflect() protoreflect.Message { + mi := &file_process_proto_msgTypes[15] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use StreamInputResponse.ProtoReflect.Descriptor instead. +func (*StreamInputResponse) Descriptor() ([]byte, []int) { + return file_process_proto_rawDescGZIP(), []int{15} +} + +type SendSignalRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + Process *ProcessSelector `protobuf:"bytes,1,opt,name=process,proto3" json:"process,omitempty"` + Signal Signal `protobuf:"varint,2,opt,name=signal,proto3,enum=process.Signal" json:"signal,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *SendSignalRequest) Reset() { + *x = SendSignalRequest{} + mi := &file_process_proto_msgTypes[16] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *SendSignalRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*SendSignalRequest) ProtoMessage() {} + +func (x *SendSignalRequest) ProtoReflect() protoreflect.Message { + mi := &file_process_proto_msgTypes[16] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use SendSignalRequest.ProtoReflect.Descriptor instead. +func (*SendSignalRequest) Descriptor() ([]byte, []int) { + return file_process_proto_rawDescGZIP(), []int{16} +} + +func (x *SendSignalRequest) GetProcess() *ProcessSelector { + if x != nil { + return x.Process + } + return nil +} + +func (x *SendSignalRequest) GetSignal() Signal { + if x != nil { + return x.Signal + } + return Signal_SIGNAL_UNSPECIFIED +} + +type SendSignalResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *SendSignalResponse) Reset() { + *x = SendSignalResponse{} + mi := &file_process_proto_msgTypes[17] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *SendSignalResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*SendSignalResponse) ProtoMessage() {} + +func (x *SendSignalResponse) ProtoReflect() protoreflect.Message { + mi := &file_process_proto_msgTypes[17] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use SendSignalResponse.ProtoReflect.Descriptor instead. +func (*SendSignalResponse) Descriptor() ([]byte, []int) { + return file_process_proto_rawDescGZIP(), []int{17} +} + +type ConnectRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + Process *ProcessSelector `protobuf:"bytes,1,opt,name=process,proto3" json:"process,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *ConnectRequest) Reset() { + *x = ConnectRequest{} + mi := &file_process_proto_msgTypes[18] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *ConnectRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ConnectRequest) ProtoMessage() {} + +func (x *ConnectRequest) ProtoReflect() protoreflect.Message { + mi := &file_process_proto_msgTypes[18] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ConnectRequest.ProtoReflect.Descriptor instead. +func (*ConnectRequest) Descriptor() ([]byte, []int) { + return file_process_proto_rawDescGZIP(), []int{18} +} + +func (x *ConnectRequest) GetProcess() *ProcessSelector { + if x != nil { + return x.Process + } + return nil +} + +type ProcessSelector struct { + state protoimpl.MessageState `protogen:"open.v1"` + // Types that are valid to be assigned to Selector: + // + // *ProcessSelector_Pid + // *ProcessSelector_Tag + Selector isProcessSelector_Selector `protobuf_oneof:"selector"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *ProcessSelector) Reset() { + *x = ProcessSelector{} + mi := &file_process_proto_msgTypes[19] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *ProcessSelector) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ProcessSelector) ProtoMessage() {} + +func (x *ProcessSelector) ProtoReflect() protoreflect.Message { + mi := &file_process_proto_msgTypes[19] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ProcessSelector.ProtoReflect.Descriptor instead. +func (*ProcessSelector) Descriptor() ([]byte, []int) { + return file_process_proto_rawDescGZIP(), []int{19} +} + +func (x *ProcessSelector) GetSelector() isProcessSelector_Selector { + if x != nil { + return x.Selector + } + return nil +} + +func (x *ProcessSelector) GetPid() uint32 { + if x != nil { + if x, ok := x.Selector.(*ProcessSelector_Pid); ok { + return x.Pid + } + } + return 0 +} + +func (x *ProcessSelector) GetTag() string { + if x != nil { + if x, ok := x.Selector.(*ProcessSelector_Tag); ok { + return x.Tag + } + } + return "" +} + +type isProcessSelector_Selector interface { + isProcessSelector_Selector() +} + +type ProcessSelector_Pid struct { + Pid uint32 `protobuf:"varint,1,opt,name=pid,proto3,oneof"` +} + +type ProcessSelector_Tag struct { + Tag string `protobuf:"bytes,2,opt,name=tag,proto3,oneof"` +} + +func (*ProcessSelector_Pid) isProcessSelector_Selector() {} + +func (*ProcessSelector_Tag) isProcessSelector_Selector() {} + +type PTY_Size struct { + state protoimpl.MessageState `protogen:"open.v1"` + Cols uint32 `protobuf:"varint,1,opt,name=cols,proto3" json:"cols,omitempty"` + Rows uint32 `protobuf:"varint,2,opt,name=rows,proto3" json:"rows,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *PTY_Size) Reset() { + *x = PTY_Size{} + mi := &file_process_proto_msgTypes[20] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *PTY_Size) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*PTY_Size) ProtoMessage() {} + +func (x *PTY_Size) ProtoReflect() protoreflect.Message { + mi := &file_process_proto_msgTypes[20] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use PTY_Size.ProtoReflect.Descriptor instead. +func (*PTY_Size) Descriptor() ([]byte, []int) { + return file_process_proto_rawDescGZIP(), []int{0, 0} +} + +func (x *PTY_Size) GetCols() uint32 { + if x != nil { + return x.Cols + } + return 0 +} + +func (x *PTY_Size) GetRows() uint32 { + if x != nil { + return x.Rows + } + return 0 +} + +type ProcessEvent_StartEvent struct { + state protoimpl.MessageState `protogen:"open.v1"` + Pid uint32 `protobuf:"varint,1,opt,name=pid,proto3" json:"pid,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *ProcessEvent_StartEvent) Reset() { + *x = ProcessEvent_StartEvent{} + mi := &file_process_proto_msgTypes[22] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *ProcessEvent_StartEvent) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ProcessEvent_StartEvent) ProtoMessage() {} + +func (x *ProcessEvent_StartEvent) ProtoReflect() protoreflect.Message { + mi := &file_process_proto_msgTypes[22] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ProcessEvent_StartEvent.ProtoReflect.Descriptor instead. +func (*ProcessEvent_StartEvent) Descriptor() ([]byte, []int) { + return file_process_proto_rawDescGZIP(), []int{8, 0} +} + +func (x *ProcessEvent_StartEvent) GetPid() uint32 { + if x != nil { + return x.Pid + } + return 0 +} + +type ProcessEvent_DataEvent struct { + state protoimpl.MessageState `protogen:"open.v1"` + // Types that are valid to be assigned to Output: + // + // *ProcessEvent_DataEvent_Stdout + // *ProcessEvent_DataEvent_Stderr + // *ProcessEvent_DataEvent_Pty + Output isProcessEvent_DataEvent_Output `protobuf_oneof:"output"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *ProcessEvent_DataEvent) Reset() { + *x = ProcessEvent_DataEvent{} + mi := &file_process_proto_msgTypes[23] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *ProcessEvent_DataEvent) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ProcessEvent_DataEvent) ProtoMessage() {} + +func (x *ProcessEvent_DataEvent) ProtoReflect() protoreflect.Message { + mi := &file_process_proto_msgTypes[23] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ProcessEvent_DataEvent.ProtoReflect.Descriptor instead. +func (*ProcessEvent_DataEvent) Descriptor() ([]byte, []int) { + return file_process_proto_rawDescGZIP(), []int{8, 1} +} + +func (x *ProcessEvent_DataEvent) GetOutput() isProcessEvent_DataEvent_Output { + if x != nil { + return x.Output + } + return nil +} + +func (x *ProcessEvent_DataEvent) GetStdout() []byte { + if x != nil { + if x, ok := x.Output.(*ProcessEvent_DataEvent_Stdout); ok { + return x.Stdout + } + } + return nil +} + +func (x *ProcessEvent_DataEvent) GetStderr() []byte { + if x != nil { + if x, ok := x.Output.(*ProcessEvent_DataEvent_Stderr); ok { + return x.Stderr + } + } + return nil +} + +func (x *ProcessEvent_DataEvent) GetPty() []byte { + if x != nil { + if x, ok := x.Output.(*ProcessEvent_DataEvent_Pty); ok { + return x.Pty + } + } + return nil +} + +type isProcessEvent_DataEvent_Output interface { + isProcessEvent_DataEvent_Output() +} + +type ProcessEvent_DataEvent_Stdout struct { + Stdout []byte `protobuf:"bytes,1,opt,name=stdout,proto3,oneof"` +} + +type ProcessEvent_DataEvent_Stderr struct { + Stderr []byte `protobuf:"bytes,2,opt,name=stderr,proto3,oneof"` +} + +type ProcessEvent_DataEvent_Pty struct { + Pty []byte `protobuf:"bytes,3,opt,name=pty,proto3,oneof"` +} + +func (*ProcessEvent_DataEvent_Stdout) isProcessEvent_DataEvent_Output() {} + +func (*ProcessEvent_DataEvent_Stderr) isProcessEvent_DataEvent_Output() {} + +func (*ProcessEvent_DataEvent_Pty) isProcessEvent_DataEvent_Output() {} + +type ProcessEvent_EndEvent struct { + state protoimpl.MessageState `protogen:"open.v1"` + ExitCode int32 `protobuf:"zigzag32,1,opt,name=exit_code,json=exitCode,proto3" json:"exit_code,omitempty"` + Exited bool `protobuf:"varint,2,opt,name=exited,proto3" json:"exited,omitempty"` + Status string `protobuf:"bytes,3,opt,name=status,proto3" json:"status,omitempty"` + Error *string `protobuf:"bytes,4,opt,name=error,proto3,oneof" json:"error,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *ProcessEvent_EndEvent) Reset() { + *x = ProcessEvent_EndEvent{} + mi := &file_process_proto_msgTypes[24] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *ProcessEvent_EndEvent) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ProcessEvent_EndEvent) ProtoMessage() {} + +func (x *ProcessEvent_EndEvent) ProtoReflect() protoreflect.Message { + mi := &file_process_proto_msgTypes[24] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ProcessEvent_EndEvent.ProtoReflect.Descriptor instead. +func (*ProcessEvent_EndEvent) Descriptor() ([]byte, []int) { + return file_process_proto_rawDescGZIP(), []int{8, 2} +} + +func (x *ProcessEvent_EndEvent) GetExitCode() int32 { + if x != nil { + return x.ExitCode + } + return 0 +} + +func (x *ProcessEvent_EndEvent) GetExited() bool { + if x != nil { + return x.Exited + } + return false +} + +func (x *ProcessEvent_EndEvent) GetStatus() string { + if x != nil { + return x.Status + } + return "" +} + +func (x *ProcessEvent_EndEvent) GetError() string { + if x != nil && x.Error != nil { + return *x.Error + } + return "" +} + +type ProcessEvent_KeepAlive struct { + state protoimpl.MessageState `protogen:"open.v1"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *ProcessEvent_KeepAlive) Reset() { + *x = ProcessEvent_KeepAlive{} + mi := &file_process_proto_msgTypes[25] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *ProcessEvent_KeepAlive) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ProcessEvent_KeepAlive) ProtoMessage() {} + +func (x *ProcessEvent_KeepAlive) ProtoReflect() protoreflect.Message { + mi := &file_process_proto_msgTypes[25] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ProcessEvent_KeepAlive.ProtoReflect.Descriptor instead. +func (*ProcessEvent_KeepAlive) Descriptor() ([]byte, []int) { + return file_process_proto_rawDescGZIP(), []int{8, 3} +} + +type StreamInputRequest_StartEvent struct { + state protoimpl.MessageState `protogen:"open.v1"` + Process *ProcessSelector `protobuf:"bytes,1,opt,name=process,proto3" json:"process,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *StreamInputRequest_StartEvent) Reset() { + *x = StreamInputRequest_StartEvent{} + mi := &file_process_proto_msgTypes[26] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *StreamInputRequest_StartEvent) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*StreamInputRequest_StartEvent) ProtoMessage() {} + +func (x *StreamInputRequest_StartEvent) ProtoReflect() protoreflect.Message { + mi := &file_process_proto_msgTypes[26] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use StreamInputRequest_StartEvent.ProtoReflect.Descriptor instead. +func (*StreamInputRequest_StartEvent) Descriptor() ([]byte, []int) { + return file_process_proto_rawDescGZIP(), []int{14, 0} +} + +func (x *StreamInputRequest_StartEvent) GetProcess() *ProcessSelector { + if x != nil { + return x.Process + } + return nil +} + +type StreamInputRequest_DataEvent struct { + state protoimpl.MessageState `protogen:"open.v1"` + Input *ProcessInput `protobuf:"bytes,2,opt,name=input,proto3" json:"input,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *StreamInputRequest_DataEvent) Reset() { + *x = StreamInputRequest_DataEvent{} + mi := &file_process_proto_msgTypes[27] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *StreamInputRequest_DataEvent) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*StreamInputRequest_DataEvent) ProtoMessage() {} + +func (x *StreamInputRequest_DataEvent) ProtoReflect() protoreflect.Message { + mi := &file_process_proto_msgTypes[27] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use StreamInputRequest_DataEvent.ProtoReflect.Descriptor instead. +func (*StreamInputRequest_DataEvent) Descriptor() ([]byte, []int) { + return file_process_proto_rawDescGZIP(), []int{14, 1} +} + +func (x *StreamInputRequest_DataEvent) GetInput() *ProcessInput { + if x != nil { + return x.Input + } + return nil +} + +type StreamInputRequest_KeepAlive struct { + state protoimpl.MessageState `protogen:"open.v1"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *StreamInputRequest_KeepAlive) Reset() { + *x = StreamInputRequest_KeepAlive{} + mi := &file_process_proto_msgTypes[28] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *StreamInputRequest_KeepAlive) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*StreamInputRequest_KeepAlive) ProtoMessage() {} + +func (x *StreamInputRequest_KeepAlive) ProtoReflect() protoreflect.Message { + mi := &file_process_proto_msgTypes[28] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use StreamInputRequest_KeepAlive.ProtoReflect.Descriptor instead. +func (*StreamInputRequest_KeepAlive) Descriptor() ([]byte, []int) { + return file_process_proto_rawDescGZIP(), []int{14, 2} +} + +var File_process_proto protoreflect.FileDescriptor + +const file_process_proto_rawDesc = "" + + "\n" + + "\rprocess.proto\x12\aprocess\"\\\n" + + "\x03PTY\x12%\n" + + "\x04size\x18\x01 \x01(\v2\x11.process.PTY.SizeR\x04size\x1a.\n" + + "\x04Size\x12\x12\n" + + "\x04cols\x18\x01 \x01(\rR\x04cols\x12\x12\n" + + "\x04rows\x18\x02 \x01(\rR\x04rows\"\xc3\x01\n" + + "\rProcessConfig\x12\x10\n" + + "\x03cmd\x18\x01 \x01(\tR\x03cmd\x12\x12\n" + + "\x04args\x18\x02 \x03(\tR\x04args\x124\n" + + "\x04envs\x18\x03 \x03(\v2 .process.ProcessConfig.EnvsEntryR\x04envs\x12\x15\n" + + "\x03cwd\x18\x04 \x01(\tH\x00R\x03cwd\x88\x01\x01\x1a7\n" + + "\tEnvsEntry\x12\x10\n" + + "\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" + + "\x05value\x18\x02 \x01(\tR\x05value:\x028\x01B\x06\n" + + "\x04_cwd\"\r\n" + + "\vListRequest\"n\n" + + "\vProcessInfo\x12.\n" + + "\x06config\x18\x01 \x01(\v2\x16.process.ProcessConfigR\x06config\x12\x10\n" + + "\x03pid\x18\x02 \x01(\rR\x03pid\x12\x15\n" + + "\x03tag\x18\x03 \x01(\tH\x00R\x03tag\x88\x01\x01B\x06\n" + + "\x04_tag\"B\n" + + "\fListResponse\x122\n" + + "\tprocesses\x18\x01 \x03(\v2\x14.process.ProcessInfoR\tprocesses\"\xb1\x01\n" + + "\fStartRequest\x120\n" + + "\aprocess\x18\x01 \x01(\v2\x16.process.ProcessConfigR\aprocess\x12#\n" + + "\x03pty\x18\x02 \x01(\v2\f.process.PTYH\x00R\x03pty\x88\x01\x01\x12\x15\n" + + "\x03tag\x18\x03 \x01(\tH\x01R\x03tag\x88\x01\x01\x12\x19\n" + + "\x05stdin\x18\x04 \x01(\bH\x02R\x05stdin\x88\x01\x01B\x06\n" + + "\x04_ptyB\x06\n" + + "\x04_tagB\b\n" + + "\x06_stdin\"p\n" + + "\rUpdateRequest\x122\n" + + "\aprocess\x18\x01 \x01(\v2\x18.process.ProcessSelectorR\aprocess\x12#\n" + + "\x03pty\x18\x02 \x01(\v2\f.process.PTYH\x00R\x03pty\x88\x01\x01B\x06\n" + + "\x04_pty\"\x10\n" + + "\x0eUpdateResponse\"\x87\x04\n" + + "\fProcessEvent\x128\n" + + "\x05start\x18\x01 \x01(\v2 .process.ProcessEvent.StartEventH\x00R\x05start\x125\n" + + "\x04data\x18\x02 \x01(\v2\x1f.process.ProcessEvent.DataEventH\x00R\x04data\x122\n" + + "\x03end\x18\x03 \x01(\v2\x1e.process.ProcessEvent.EndEventH\x00R\x03end\x12?\n" + + "\tkeepalive\x18\x04 \x01(\v2\x1f.process.ProcessEvent.KeepAliveH\x00R\tkeepalive\x1a\x1e\n" + + "\n" + + "StartEvent\x12\x10\n" + + "\x03pid\x18\x01 \x01(\rR\x03pid\x1a]\n" + + "\tDataEvent\x12\x18\n" + + "\x06stdout\x18\x01 \x01(\fH\x00R\x06stdout\x12\x18\n" + + "\x06stderr\x18\x02 \x01(\fH\x00R\x06stderr\x12\x12\n" + + "\x03pty\x18\x03 \x01(\fH\x00R\x03ptyB\b\n" + + "\x06output\x1a|\n" + + "\bEndEvent\x12\x1b\n" + + "\texit_code\x18\x01 \x01(\x11R\bexitCode\x12\x16\n" + + "\x06exited\x18\x02 \x01(\bR\x06exited\x12\x16\n" + + "\x06status\x18\x03 \x01(\tR\x06status\x12\x19\n" + + "\x05error\x18\x04 \x01(\tH\x00R\x05error\x88\x01\x01B\b\n" + + "\x06_error\x1a\v\n" + + "\tKeepAliveB\a\n" + + "\x05event\"<\n" + + "\rStartResponse\x12+\n" + + "\x05event\x18\x01 \x01(\v2\x15.process.ProcessEventR\x05event\">\n" + + "\x0fConnectResponse\x12+\n" + + "\x05event\x18\x01 \x01(\v2\x15.process.ProcessEventR\x05event\"s\n" + + "\x10SendInputRequest\x122\n" + + "\aprocess\x18\x01 \x01(\v2\x18.process.ProcessSelectorR\aprocess\x12+\n" + + "\x05input\x18\x02 \x01(\v2\x15.process.ProcessInputR\x05input\"\x13\n" + + "\x11SendInputResponse\"C\n" + + "\fProcessInput\x12\x16\n" + + "\x05stdin\x18\x01 \x01(\fH\x00R\x05stdin\x12\x12\n" + + "\x03pty\x18\x02 \x01(\fH\x00R\x03ptyB\a\n" + + "\x05input\"\xea\x02\n" + + "\x12StreamInputRequest\x12>\n" + + "\x05start\x18\x01 \x01(\v2&.process.StreamInputRequest.StartEventH\x00R\x05start\x12;\n" + + "\x04data\x18\x02 \x01(\v2%.process.StreamInputRequest.DataEventH\x00R\x04data\x12E\n" + + "\tkeepalive\x18\x03 \x01(\v2%.process.StreamInputRequest.KeepAliveH\x00R\tkeepalive\x1a@\n" + + "\n" + + "StartEvent\x122\n" + + "\aprocess\x18\x01 \x01(\v2\x18.process.ProcessSelectorR\aprocess\x1a8\n" + + "\tDataEvent\x12+\n" + + "\x05input\x18\x02 \x01(\v2\x15.process.ProcessInputR\x05input\x1a\v\n" + + "\tKeepAliveB\a\n" + + "\x05event\"\x15\n" + + "\x13StreamInputResponse\"p\n" + + "\x11SendSignalRequest\x122\n" + + "\aprocess\x18\x01 \x01(\v2\x18.process.ProcessSelectorR\aprocess\x12'\n" + + "\x06signal\x18\x02 \x01(\x0e2\x0f.process.SignalR\x06signal\"\x14\n" + + "\x12SendSignalResponse\"D\n" + + "\x0eConnectRequest\x122\n" + + "\aprocess\x18\x01 \x01(\v2\x18.process.ProcessSelectorR\aprocess\"E\n" + + "\x0fProcessSelector\x12\x12\n" + + "\x03pid\x18\x01 \x01(\rH\x00R\x03pid\x12\x12\n" + + "\x03tag\x18\x02 \x01(\tH\x00R\x03tagB\n" + + "\n" + + "\bselector*H\n" + + "\x06Signal\x12\x16\n" + + "\x12SIGNAL_UNSPECIFIED\x10\x00\x12\x12\n" + + "\x0eSIGNAL_SIGTERM\x10\x0f\x12\x12\n" + + "\x0eSIGNAL_SIGKILL\x10\t2\xca\x03\n" + + "\aProcess\x123\n" + + "\x04List\x12\x14.process.ListRequest\x1a\x15.process.ListResponse\x12>\n" + + "\aConnect\x12\x17.process.ConnectRequest\x1a\x18.process.ConnectResponse0\x01\x128\n" + + "\x05Start\x12\x15.process.StartRequest\x1a\x16.process.StartResponse0\x01\x129\n" + + "\x06Update\x12\x16.process.UpdateRequest\x1a\x17.process.UpdateResponse\x12J\n" + + "\vStreamInput\x12\x1b.process.StreamInputRequest\x1a\x1c.process.StreamInputResponse(\x01\x12B\n" + + "\tSendInput\x12\x19.process.SendInputRequest\x1a\x1a.process.SendInputResponse\x12E\n" + + "\n" + + "SendSignal\x12\x1a.process.SendSignalRequest\x1a\x1b.process.SendSignalResponseB3Z1station/pkg/harness/sandbox/e2b/proto/gen/processb\x06proto3" + +var ( + file_process_proto_rawDescOnce sync.Once + file_process_proto_rawDescData []byte +) + +func file_process_proto_rawDescGZIP() []byte { + file_process_proto_rawDescOnce.Do(func() { + file_process_proto_rawDescData = protoimpl.X.CompressGZIP(unsafe.Slice(unsafe.StringData(file_process_proto_rawDesc), len(file_process_proto_rawDesc))) + }) + return file_process_proto_rawDescData +} + +var file_process_proto_enumTypes = make([]protoimpl.EnumInfo, 1) +var file_process_proto_msgTypes = make([]protoimpl.MessageInfo, 29) +var file_process_proto_goTypes = []any{ + (Signal)(0), // 0: process.Signal + (*PTY)(nil), // 1: process.PTY + (*ProcessConfig)(nil), // 2: process.ProcessConfig + (*ListRequest)(nil), // 3: process.ListRequest + (*ProcessInfo)(nil), // 4: process.ProcessInfo + (*ListResponse)(nil), // 5: process.ListResponse + (*StartRequest)(nil), // 6: process.StartRequest + (*UpdateRequest)(nil), // 7: process.UpdateRequest + (*UpdateResponse)(nil), // 8: process.UpdateResponse + (*ProcessEvent)(nil), // 9: process.ProcessEvent + (*StartResponse)(nil), // 10: process.StartResponse + (*ConnectResponse)(nil), // 11: process.ConnectResponse + (*SendInputRequest)(nil), // 12: process.SendInputRequest + (*SendInputResponse)(nil), // 13: process.SendInputResponse + (*ProcessInput)(nil), // 14: process.ProcessInput + (*StreamInputRequest)(nil), // 15: process.StreamInputRequest + (*StreamInputResponse)(nil), // 16: process.StreamInputResponse + (*SendSignalRequest)(nil), // 17: process.SendSignalRequest + (*SendSignalResponse)(nil), // 18: process.SendSignalResponse + (*ConnectRequest)(nil), // 19: process.ConnectRequest + (*ProcessSelector)(nil), // 20: process.ProcessSelector + (*PTY_Size)(nil), // 21: process.PTY.Size + nil, // 22: process.ProcessConfig.EnvsEntry + (*ProcessEvent_StartEvent)(nil), // 23: process.ProcessEvent.StartEvent + (*ProcessEvent_DataEvent)(nil), // 24: process.ProcessEvent.DataEvent + (*ProcessEvent_EndEvent)(nil), // 25: process.ProcessEvent.EndEvent + (*ProcessEvent_KeepAlive)(nil), // 26: process.ProcessEvent.KeepAlive + (*StreamInputRequest_StartEvent)(nil), // 27: process.StreamInputRequest.StartEvent + (*StreamInputRequest_DataEvent)(nil), // 28: process.StreamInputRequest.DataEvent + (*StreamInputRequest_KeepAlive)(nil), // 29: process.StreamInputRequest.KeepAlive +} +var file_process_proto_depIdxs = []int32{ + 21, // 0: process.PTY.size:type_name -> process.PTY.Size + 22, // 1: process.ProcessConfig.envs:type_name -> process.ProcessConfig.EnvsEntry + 2, // 2: process.ProcessInfo.config:type_name -> process.ProcessConfig + 4, // 3: process.ListResponse.processes:type_name -> process.ProcessInfo + 2, // 4: process.StartRequest.process:type_name -> process.ProcessConfig + 1, // 5: process.StartRequest.pty:type_name -> process.PTY + 20, // 6: process.UpdateRequest.process:type_name -> process.ProcessSelector + 1, // 7: process.UpdateRequest.pty:type_name -> process.PTY + 23, // 8: process.ProcessEvent.start:type_name -> process.ProcessEvent.StartEvent + 24, // 9: process.ProcessEvent.data:type_name -> process.ProcessEvent.DataEvent + 25, // 10: process.ProcessEvent.end:type_name -> process.ProcessEvent.EndEvent + 26, // 11: process.ProcessEvent.keepalive:type_name -> process.ProcessEvent.KeepAlive + 9, // 12: process.StartResponse.event:type_name -> process.ProcessEvent + 9, // 13: process.ConnectResponse.event:type_name -> process.ProcessEvent + 20, // 14: process.SendInputRequest.process:type_name -> process.ProcessSelector + 14, // 15: process.SendInputRequest.input:type_name -> process.ProcessInput + 27, // 16: process.StreamInputRequest.start:type_name -> process.StreamInputRequest.StartEvent + 28, // 17: process.StreamInputRequest.data:type_name -> process.StreamInputRequest.DataEvent + 29, // 18: process.StreamInputRequest.keepalive:type_name -> process.StreamInputRequest.KeepAlive + 20, // 19: process.SendSignalRequest.process:type_name -> process.ProcessSelector + 0, // 20: process.SendSignalRequest.signal:type_name -> process.Signal + 20, // 21: process.ConnectRequest.process:type_name -> process.ProcessSelector + 20, // 22: process.StreamInputRequest.StartEvent.process:type_name -> process.ProcessSelector + 14, // 23: process.StreamInputRequest.DataEvent.input:type_name -> process.ProcessInput + 3, // 24: process.Process.List:input_type -> process.ListRequest + 19, // 25: process.Process.Connect:input_type -> process.ConnectRequest + 6, // 26: process.Process.Start:input_type -> process.StartRequest + 7, // 27: process.Process.Update:input_type -> process.UpdateRequest + 15, // 28: process.Process.StreamInput:input_type -> process.StreamInputRequest + 12, // 29: process.Process.SendInput:input_type -> process.SendInputRequest + 17, // 30: process.Process.SendSignal:input_type -> process.SendSignalRequest + 5, // 31: process.Process.List:output_type -> process.ListResponse + 11, // 32: process.Process.Connect:output_type -> process.ConnectResponse + 10, // 33: process.Process.Start:output_type -> process.StartResponse + 8, // 34: process.Process.Update:output_type -> process.UpdateResponse + 16, // 35: process.Process.StreamInput:output_type -> process.StreamInputResponse + 13, // 36: process.Process.SendInput:output_type -> process.SendInputResponse + 18, // 37: process.Process.SendSignal:output_type -> process.SendSignalResponse + 31, // [31:38] is the sub-list for method output_type + 24, // [24:31] is the sub-list for method input_type + 24, // [24:24] is the sub-list for extension type_name + 24, // [24:24] is the sub-list for extension extendee + 0, // [0:24] is the sub-list for field type_name +} + +func init() { file_process_proto_init() } +func file_process_proto_init() { + if File_process_proto != nil { + return + } + file_process_proto_msgTypes[1].OneofWrappers = []any{} + file_process_proto_msgTypes[3].OneofWrappers = []any{} + file_process_proto_msgTypes[5].OneofWrappers = []any{} + file_process_proto_msgTypes[6].OneofWrappers = []any{} + file_process_proto_msgTypes[8].OneofWrappers = []any{ + (*ProcessEvent_Start)(nil), + (*ProcessEvent_Data)(nil), + (*ProcessEvent_End)(nil), + (*ProcessEvent_Keepalive)(nil), + } + file_process_proto_msgTypes[13].OneofWrappers = []any{ + (*ProcessInput_Stdin)(nil), + (*ProcessInput_Pty)(nil), + } + file_process_proto_msgTypes[14].OneofWrappers = []any{ + (*StreamInputRequest_Start)(nil), + (*StreamInputRequest_Data)(nil), + (*StreamInputRequest_Keepalive)(nil), + } + file_process_proto_msgTypes[19].OneofWrappers = []any{ + (*ProcessSelector_Pid)(nil), + (*ProcessSelector_Tag)(nil), + } + file_process_proto_msgTypes[23].OneofWrappers = []any{ + (*ProcessEvent_DataEvent_Stdout)(nil), + (*ProcessEvent_DataEvent_Stderr)(nil), + (*ProcessEvent_DataEvent_Pty)(nil), + } + file_process_proto_msgTypes[24].OneofWrappers = []any{} + type x struct{} + out := protoimpl.TypeBuilder{ + File: protoimpl.DescBuilder{ + GoPackagePath: reflect.TypeOf(x{}).PkgPath(), + RawDescriptor: unsafe.Slice(unsafe.StringData(file_process_proto_rawDesc), len(file_process_proto_rawDesc)), + NumEnums: 1, + NumMessages: 29, + NumExtensions: 0, + NumServices: 1, + }, + GoTypes: file_process_proto_goTypes, + DependencyIndexes: file_process_proto_depIdxs, + EnumInfos: file_process_proto_enumTypes, + MessageInfos: file_process_proto_msgTypes, + }.Build() + File_process_proto = out.File + file_process_proto_goTypes = nil + file_process_proto_depIdxs = nil +} diff --git a/pkg/harness/sandbox/e2b/proto/gen/processconnect/process.connect.go b/pkg/harness/sandbox/e2b/proto/gen/processconnect/process.connect.go new file mode 100644 index 00000000..7ca9f832 --- /dev/null +++ b/pkg/harness/sandbox/e2b/proto/gen/processconnect/process.connect.go @@ -0,0 +1,278 @@ +// Code generated by protoc-gen-connect-go. DO NOT EDIT. +// +// Source: process.proto + +package processconnect + +import ( + connect "connectrpc.com/connect" + context "context" + errors "errors" + http "net/http" + process "station/pkg/harness/sandbox/e2b/proto/gen/process" + strings "strings" +) + +// This is a compile-time assertion to ensure that this generated file and the connect package are +// compatible. If you get a compiler error that this constant is not defined, this code was +// generated with a version of connect newer than the one compiled into your binary. You can fix the +// problem by either regenerating this code with an older version of connect or updating the connect +// version compiled into your binary. +const _ = connect.IsAtLeastVersion1_13_0 + +const ( + // ProcessName is the fully-qualified name of the Process service. + ProcessName = "process.Process" +) + +// These constants are the fully-qualified names of the RPCs defined in this package. They're +// exposed at runtime as Spec.Procedure and as the final two segments of the HTTP route. +// +// Note that these are different from the fully-qualified method names used by +// google.golang.org/protobuf/reflect/protoreflect. To convert from these constants to +// reflection-formatted method names, remove the leading slash and convert the remaining slash to a +// period. +const ( + // ProcessListProcedure is the fully-qualified name of the Process's List RPC. + ProcessListProcedure = "/process.Process/List" + // ProcessConnectProcedure is the fully-qualified name of the Process's Connect RPC. + ProcessConnectProcedure = "/process.Process/Connect" + // ProcessStartProcedure is the fully-qualified name of the Process's Start RPC. + ProcessStartProcedure = "/process.Process/Start" + // ProcessUpdateProcedure is the fully-qualified name of the Process's Update RPC. + ProcessUpdateProcedure = "/process.Process/Update" + // ProcessStreamInputProcedure is the fully-qualified name of the Process's StreamInput RPC. + ProcessStreamInputProcedure = "/process.Process/StreamInput" + // ProcessSendInputProcedure is the fully-qualified name of the Process's SendInput RPC. + ProcessSendInputProcedure = "/process.Process/SendInput" + // ProcessSendSignalProcedure is the fully-qualified name of the Process's SendSignal RPC. + ProcessSendSignalProcedure = "/process.Process/SendSignal" +) + +// ProcessClient is a client for the process.Process service. +type ProcessClient interface { + List(context.Context, *connect.Request[process.ListRequest]) (*connect.Response[process.ListResponse], error) + Connect(context.Context, *connect.Request[process.ConnectRequest]) (*connect.ServerStreamForClient[process.ConnectResponse], error) + Start(context.Context, *connect.Request[process.StartRequest]) (*connect.ServerStreamForClient[process.StartResponse], error) + Update(context.Context, *connect.Request[process.UpdateRequest]) (*connect.Response[process.UpdateResponse], error) + // Client input stream ensures ordering of messages + StreamInput(context.Context) *connect.ClientStreamForClient[process.StreamInputRequest, process.StreamInputResponse] + SendInput(context.Context, *connect.Request[process.SendInputRequest]) (*connect.Response[process.SendInputResponse], error) + SendSignal(context.Context, *connect.Request[process.SendSignalRequest]) (*connect.Response[process.SendSignalResponse], error) +} + +// NewProcessClient constructs a client for the process.Process service. By default, it uses the +// Connect protocol with the binary Protobuf Codec, asks for gzipped responses, and sends +// uncompressed requests. To use the gRPC or gRPC-Web protocols, supply the connect.WithGRPC() or +// connect.WithGRPCWeb() options. +// +// The URL supplied here should be the base URL for the Connect or gRPC server (for example, +// http://api.acme.com or https://acme.com/grpc). +func NewProcessClient(httpClient connect.HTTPClient, baseURL string, opts ...connect.ClientOption) ProcessClient { + baseURL = strings.TrimRight(baseURL, "/") + processMethods := process.File_process_proto.Services().ByName("Process").Methods() + return &processClient{ + list: connect.NewClient[process.ListRequest, process.ListResponse]( + httpClient, + baseURL+ProcessListProcedure, + connect.WithSchema(processMethods.ByName("List")), + connect.WithClientOptions(opts...), + ), + connect: connect.NewClient[process.ConnectRequest, process.ConnectResponse]( + httpClient, + baseURL+ProcessConnectProcedure, + connect.WithSchema(processMethods.ByName("Connect")), + connect.WithClientOptions(opts...), + ), + start: connect.NewClient[process.StartRequest, process.StartResponse]( + httpClient, + baseURL+ProcessStartProcedure, + connect.WithSchema(processMethods.ByName("Start")), + connect.WithClientOptions(opts...), + ), + update: connect.NewClient[process.UpdateRequest, process.UpdateResponse]( + httpClient, + baseURL+ProcessUpdateProcedure, + connect.WithSchema(processMethods.ByName("Update")), + connect.WithClientOptions(opts...), + ), + streamInput: connect.NewClient[process.StreamInputRequest, process.StreamInputResponse]( + httpClient, + baseURL+ProcessStreamInputProcedure, + connect.WithSchema(processMethods.ByName("StreamInput")), + connect.WithClientOptions(opts...), + ), + sendInput: connect.NewClient[process.SendInputRequest, process.SendInputResponse]( + httpClient, + baseURL+ProcessSendInputProcedure, + connect.WithSchema(processMethods.ByName("SendInput")), + connect.WithClientOptions(opts...), + ), + sendSignal: connect.NewClient[process.SendSignalRequest, process.SendSignalResponse]( + httpClient, + baseURL+ProcessSendSignalProcedure, + connect.WithSchema(processMethods.ByName("SendSignal")), + connect.WithClientOptions(opts...), + ), + } +} + +// processClient implements ProcessClient. +type processClient struct { + list *connect.Client[process.ListRequest, process.ListResponse] + connect *connect.Client[process.ConnectRequest, process.ConnectResponse] + start *connect.Client[process.StartRequest, process.StartResponse] + update *connect.Client[process.UpdateRequest, process.UpdateResponse] + streamInput *connect.Client[process.StreamInputRequest, process.StreamInputResponse] + sendInput *connect.Client[process.SendInputRequest, process.SendInputResponse] + sendSignal *connect.Client[process.SendSignalRequest, process.SendSignalResponse] +} + +// List calls process.Process.List. +func (c *processClient) List(ctx context.Context, req *connect.Request[process.ListRequest]) (*connect.Response[process.ListResponse], error) { + return c.list.CallUnary(ctx, req) +} + +// Connect calls process.Process.Connect. +func (c *processClient) Connect(ctx context.Context, req *connect.Request[process.ConnectRequest]) (*connect.ServerStreamForClient[process.ConnectResponse], error) { + return c.connect.CallServerStream(ctx, req) +} + +// Start calls process.Process.Start. +func (c *processClient) Start(ctx context.Context, req *connect.Request[process.StartRequest]) (*connect.ServerStreamForClient[process.StartResponse], error) { + return c.start.CallServerStream(ctx, req) +} + +// Update calls process.Process.Update. +func (c *processClient) Update(ctx context.Context, req *connect.Request[process.UpdateRequest]) (*connect.Response[process.UpdateResponse], error) { + return c.update.CallUnary(ctx, req) +} + +// StreamInput calls process.Process.StreamInput. +func (c *processClient) StreamInput(ctx context.Context) *connect.ClientStreamForClient[process.StreamInputRequest, process.StreamInputResponse] { + return c.streamInput.CallClientStream(ctx) +} + +// SendInput calls process.Process.SendInput. +func (c *processClient) SendInput(ctx context.Context, req *connect.Request[process.SendInputRequest]) (*connect.Response[process.SendInputResponse], error) { + return c.sendInput.CallUnary(ctx, req) +} + +// SendSignal calls process.Process.SendSignal. +func (c *processClient) SendSignal(ctx context.Context, req *connect.Request[process.SendSignalRequest]) (*connect.Response[process.SendSignalResponse], error) { + return c.sendSignal.CallUnary(ctx, req) +} + +// ProcessHandler is an implementation of the process.Process service. +type ProcessHandler interface { + List(context.Context, *connect.Request[process.ListRequest]) (*connect.Response[process.ListResponse], error) + Connect(context.Context, *connect.Request[process.ConnectRequest], *connect.ServerStream[process.ConnectResponse]) error + Start(context.Context, *connect.Request[process.StartRequest], *connect.ServerStream[process.StartResponse]) error + Update(context.Context, *connect.Request[process.UpdateRequest]) (*connect.Response[process.UpdateResponse], error) + // Client input stream ensures ordering of messages + StreamInput(context.Context, *connect.ClientStream[process.StreamInputRequest]) (*connect.Response[process.StreamInputResponse], error) + SendInput(context.Context, *connect.Request[process.SendInputRequest]) (*connect.Response[process.SendInputResponse], error) + SendSignal(context.Context, *connect.Request[process.SendSignalRequest]) (*connect.Response[process.SendSignalResponse], error) +} + +// NewProcessHandler builds an HTTP handler from the service implementation. It returns the path on +// which to mount the handler and the handler itself. +// +// By default, handlers support the Connect, gRPC, and gRPC-Web protocols with the binary Protobuf +// and JSON codecs. They also support gzip compression. +func NewProcessHandler(svc ProcessHandler, opts ...connect.HandlerOption) (string, http.Handler) { + processMethods := process.File_process_proto.Services().ByName("Process").Methods() + processListHandler := connect.NewUnaryHandler( + ProcessListProcedure, + svc.List, + connect.WithSchema(processMethods.ByName("List")), + connect.WithHandlerOptions(opts...), + ) + processConnectHandler := connect.NewServerStreamHandler( + ProcessConnectProcedure, + svc.Connect, + connect.WithSchema(processMethods.ByName("Connect")), + connect.WithHandlerOptions(opts...), + ) + processStartHandler := connect.NewServerStreamHandler( + ProcessStartProcedure, + svc.Start, + connect.WithSchema(processMethods.ByName("Start")), + connect.WithHandlerOptions(opts...), + ) + processUpdateHandler := connect.NewUnaryHandler( + ProcessUpdateProcedure, + svc.Update, + connect.WithSchema(processMethods.ByName("Update")), + connect.WithHandlerOptions(opts...), + ) + processStreamInputHandler := connect.NewClientStreamHandler( + ProcessStreamInputProcedure, + svc.StreamInput, + connect.WithSchema(processMethods.ByName("StreamInput")), + connect.WithHandlerOptions(opts...), + ) + processSendInputHandler := connect.NewUnaryHandler( + ProcessSendInputProcedure, + svc.SendInput, + connect.WithSchema(processMethods.ByName("SendInput")), + connect.WithHandlerOptions(opts...), + ) + processSendSignalHandler := connect.NewUnaryHandler( + ProcessSendSignalProcedure, + svc.SendSignal, + connect.WithSchema(processMethods.ByName("SendSignal")), + connect.WithHandlerOptions(opts...), + ) + return "/process.Process/", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case ProcessListProcedure: + processListHandler.ServeHTTP(w, r) + case ProcessConnectProcedure: + processConnectHandler.ServeHTTP(w, r) + case ProcessStartProcedure: + processStartHandler.ServeHTTP(w, r) + case ProcessUpdateProcedure: + processUpdateHandler.ServeHTTP(w, r) + case ProcessStreamInputProcedure: + processStreamInputHandler.ServeHTTP(w, r) + case ProcessSendInputProcedure: + processSendInputHandler.ServeHTTP(w, r) + case ProcessSendSignalProcedure: + processSendSignalHandler.ServeHTTP(w, r) + default: + http.NotFound(w, r) + } + }) +} + +// UnimplementedProcessHandler returns CodeUnimplemented from all methods. +type UnimplementedProcessHandler struct{} + +func (UnimplementedProcessHandler) List(context.Context, *connect.Request[process.ListRequest]) (*connect.Response[process.ListResponse], error) { + return nil, connect.NewError(connect.CodeUnimplemented, errors.New("process.Process.List is not implemented")) +} + +func (UnimplementedProcessHandler) Connect(context.Context, *connect.Request[process.ConnectRequest], *connect.ServerStream[process.ConnectResponse]) error { + return connect.NewError(connect.CodeUnimplemented, errors.New("process.Process.Connect is not implemented")) +} + +func (UnimplementedProcessHandler) Start(context.Context, *connect.Request[process.StartRequest], *connect.ServerStream[process.StartResponse]) error { + return connect.NewError(connect.CodeUnimplemented, errors.New("process.Process.Start is not implemented")) +} + +func (UnimplementedProcessHandler) Update(context.Context, *connect.Request[process.UpdateRequest]) (*connect.Response[process.UpdateResponse], error) { + return nil, connect.NewError(connect.CodeUnimplemented, errors.New("process.Process.Update is not implemented")) +} + +func (UnimplementedProcessHandler) StreamInput(context.Context, *connect.ClientStream[process.StreamInputRequest]) (*connect.Response[process.StreamInputResponse], error) { + return nil, connect.NewError(connect.CodeUnimplemented, errors.New("process.Process.StreamInput is not implemented")) +} + +func (UnimplementedProcessHandler) SendInput(context.Context, *connect.Request[process.SendInputRequest]) (*connect.Response[process.SendInputResponse], error) { + return nil, connect.NewError(connect.CodeUnimplemented, errors.New("process.Process.SendInput is not implemented")) +} + +func (UnimplementedProcessHandler) SendSignal(context.Context, *connect.Request[process.SendSignalRequest]) (*connect.Response[process.SendSignalResponse], error) { + return nil, connect.NewError(connect.CodeUnimplemented, errors.New("process.Process.SendSignal is not implemented")) +} diff --git a/pkg/harness/sandbox/e2b/proto/process.proto b/pkg/harness/sandbox/e2b/proto/process.proto new file mode 100644 index 00000000..a8c4b50c --- /dev/null +++ b/pkg/harness/sandbox/e2b/proto/process.proto @@ -0,0 +1,161 @@ +syntax = "proto3"; + +package process; + +option go_package = "station/pkg/harness/sandbox/e2b/proto/gen/process"; + +service Process { + rpc List(ListRequest) returns (ListResponse); + + rpc Connect(ConnectRequest) returns (stream ConnectResponse); + rpc Start(StartRequest) returns (stream StartResponse); + + rpc Update(UpdateRequest) returns (UpdateResponse); + + // Client input stream ensures ordering of messages + rpc StreamInput(stream StreamInputRequest) returns (StreamInputResponse); + rpc SendInput(SendInputRequest) returns (SendInputResponse); + rpc SendSignal(SendSignalRequest) returns (SendSignalResponse); +} + +message PTY { + Size size = 1; + + message Size { + uint32 cols = 1; + uint32 rows = 2; + } +} + +message ProcessConfig { + string cmd = 1; + repeated string args = 2; + + map envs = 3; + optional string cwd = 4; +} + +message ListRequest {} + +message ProcessInfo { + ProcessConfig config = 1; + uint32 pid = 2; + optional string tag = 3; +} + +message ListResponse { + repeated ProcessInfo processes = 1; +} + +message StartRequest { + ProcessConfig process = 1; + optional PTY pty = 2; + optional string tag = 3; + optional bool stdin = 4; +} + +message UpdateRequest { + ProcessSelector process = 1; + + optional PTY pty = 2; +} + +message UpdateResponse {} + +message ProcessEvent { + oneof event { + StartEvent start = 1; + DataEvent data = 2; + EndEvent end = 3; + KeepAlive keepalive = 4; + } + + message StartEvent { + uint32 pid = 1; + } + + message DataEvent { + oneof output { + bytes stdout = 1; + bytes stderr = 2; + bytes pty = 3; + } + } + + message EndEvent { + sint32 exit_code = 1; + bool exited = 2; + string status = 3; + optional string error = 4; + } + + message KeepAlive {} +} + +message StartResponse { + ProcessEvent event = 1; +} + +message ConnectResponse { + ProcessEvent event = 1; +} + +message SendInputRequest { + ProcessSelector process = 1; + + ProcessInput input = 2; +} + +message SendInputResponse {} + +message ProcessInput { + oneof input { + bytes stdin = 1; + bytes pty = 2; + } +} + +message StreamInputRequest { + oneof event { + StartEvent start = 1; + DataEvent data = 2; + KeepAlive keepalive = 3; + } + + message StartEvent { + ProcessSelector process = 1; + } + + message DataEvent { + ProcessInput input = 2; + } + + message KeepAlive {} +} + +message StreamInputResponse {} + +enum Signal { + SIGNAL_UNSPECIFIED = 0; + SIGNAL_SIGTERM = 15; + SIGNAL_SIGKILL = 9; +} + +message SendSignalRequest { + ProcessSelector process = 1; + + Signal signal = 2; +} + +message SendSignalResponse {} + +message ConnectRequest { + ProcessSelector process = 1; +} + +message ProcessSelector { + oneof selector { + uint32 pid = 1; + string tag = 2; + } +} diff --git a/pkg/harness/sandbox/e2b/streaming.go b/pkg/harness/sandbox/e2b/streaming.go new file mode 100644 index 00000000..cd941de7 --- /dev/null +++ b/pkg/harness/sandbox/e2b/streaming.go @@ -0,0 +1,195 @@ +package e2b + +import ( + "context" + "fmt" + "net/http" + + "connectrpc.com/connect" + "station/pkg/harness/sandbox/e2b/proto/gen/process" + "station/pkg/harness/sandbox/e2b/proto/gen/processconnect" +) + +type StreamingClient struct { + processClient processconnect.ProcessClient + accessToken string +} + +func NewStreamingClient(envdURL, accessToken string) *StreamingClient { + httpClient := &http.Client{} + + opts := []connect.ClientOption{ + connect.WithInterceptors(&authInterceptor{token: accessToken}), + } + + return &StreamingClient{ + processClient: processconnect.NewProcessClient(httpClient, envdURL, opts...), + accessToken: accessToken, + } +} + +type authInterceptor struct { + token string +} + +func (i *authInterceptor) WrapUnary(next connect.UnaryFunc) connect.UnaryFunc { + return func(ctx context.Context, req connect.AnyRequest) (connect.AnyResponse, error) { + req.Header().Set("X-Access-Token", i.token) + return next(ctx, req) + } +} + +func (i *authInterceptor) WrapStreamingClient(next connect.StreamingClientFunc) connect.StreamingClientFunc { + return func(ctx context.Context, spec connect.Spec) connect.StreamingClientConn { + conn := next(ctx, spec) + conn.RequestHeader().Set("X-Access-Token", i.token) + return conn + } +} + +func (i *authInterceptor) WrapStreamingHandler(next connect.StreamingHandlerFunc) connect.StreamingHandlerFunc { + return next +} + +type StreamExecOptions struct { + Cmd string + Args []string + Cwd string + Env map[string]string + OnStdout func(data []byte) + OnStderr func(data []byte) + Stdin bool +} + +type StreamExecHandle struct { + stream *connect.ServerStreamForClient[process.StartResponse] + pid uint32 + onStdout func([]byte) + onStderr func([]byte) + done chan struct{} + exitCode int32 + exitError string + err error +} + +func (c *StreamingClient) ExecStream(ctx context.Context, opts StreamExecOptions) (*StreamExecHandle, error) { + envs := make(map[string]string) + for k, v := range opts.Env { + envs[k] = v + } + + req := &process.StartRequest{ + Process: &process.ProcessConfig{ + Cmd: opts.Cmd, + Args: opts.Args, + Envs: envs, + Cwd: &opts.Cwd, + }, + Stdin: &opts.Stdin, + } + + stream, err := c.processClient.Start(ctx, connect.NewRequest(req)) + if err != nil { + return nil, fmt.Errorf("start process: %w", err) + } + + handle := &StreamExecHandle{ + stream: stream, + onStdout: opts.OnStdout, + onStderr: opts.OnStderr, + done: make(chan struct{}), + } + + go handle.consumeEvents() + + return handle, nil +} + +func (h *StreamExecHandle) consumeEvents() { + defer close(h.done) + + for h.stream.Receive() { + resp := h.stream.Msg() + if resp.Event == nil { + continue + } + + switch ev := resp.Event.Event.(type) { + case *process.ProcessEvent_Start: + h.pid = ev.Start.Pid + + case *process.ProcessEvent_Data: + switch out := ev.Data.Output.(type) { + case *process.ProcessEvent_DataEvent_Stdout: + if h.onStdout != nil { + h.onStdout(out.Stdout) + } + case *process.ProcessEvent_DataEvent_Stderr: + if h.onStderr != nil { + h.onStderr(out.Stderr) + } + } + + case *process.ProcessEvent_End: + h.exitCode = ev.End.ExitCode + if ev.End.Error != nil { + h.exitError = *ev.End.Error + } + return + + case *process.ProcessEvent_Keepalive: + continue + } + } + + if err := h.stream.Err(); err != nil { + h.err = err + } +} + +func (h *StreamExecHandle) PID() uint32 { + return h.pid +} + +func (h *StreamExecHandle) Wait() (exitCode int32, exitError string, err error) { + <-h.done + return h.exitCode, h.exitError, h.err +} + +func (h *StreamExecHandle) Close() error { + return h.stream.Close() +} + +func (c *StreamingClient) ListProcesses(ctx context.Context) ([]*process.ProcessInfo, error) { + resp, err := c.processClient.List(ctx, connect.NewRequest(&process.ListRequest{})) + if err != nil { + return nil, err + } + return resp.Msg.Processes, nil +} + +func (c *StreamingClient) SendSignal(ctx context.Context, pid uint32, signal process.Signal) error { + _, err := c.processClient.SendSignal(ctx, connect.NewRequest(&process.SendSignalRequest{ + Process: &process.ProcessSelector{ + Selector: &process.ProcessSelector_Pid{Pid: pid}, + }, + Signal: signal, + })) + return err +} + +func (c *StreamingClient) Kill(ctx context.Context, pid uint32) error { + return c.SendSignal(ctx, pid, process.Signal_SIGNAL_SIGKILL) +} + +func (c *StreamingClient) SendStdin(ctx context.Context, pid uint32, data []byte) error { + _, err := c.processClient.SendInput(ctx, connect.NewRequest(&process.SendInputRequest{ + Process: &process.ProcessSelector{ + Selector: &process.ProcessSelector_Pid{Pid: pid}, + }, + Input: &process.ProcessInput{ + Input: &process.ProcessInput_Stdin{Stdin: data}, + }, + })) + return err +} diff --git a/pkg/harness/sandbox/e2b_test.go b/pkg/harness/sandbox/e2b_test.go new file mode 100644 index 00000000..697ce1d4 --- /dev/null +++ b/pkg/harness/sandbox/e2b_test.go @@ -0,0 +1,350 @@ +package sandbox + +import ( + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "os" + "strings" + "testing" + "time" +) + +func TestE2BSandbox_Create(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == "POST" && r.URL.Path == "/sandboxes" { + if r.Header.Get("X-API-KEY") != "test-api-key" { + w.WriteHeader(http.StatusUnauthorized) + return + } + + var req map[string]interface{} + json.NewDecoder(r.Body).Decode(&req) + + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusCreated) + json.NewEncoder(w).Encode(map[string]interface{}{ + "sandboxID": "sb-test-123", + "templateID": req["templateID"], + "domain": "e2b.dev", + "envdAccessToken": "test-access-token", + "envdVersion": "0.1.0", + "startedAt": time.Now().Format(time.RFC3339), + "endAt": time.Now().Add(5 * time.Minute).Format(time.RFC3339), + }) + return + } + w.WriteHeader(http.StatusNotFound) + })) + defer server.Close() + + sb, err := NewE2BSandboxWithAPIURL(Config{ + Mode: ModeE2B, + Image: "base", + Timeout: 5 * time.Minute, + }, E2BConfig{ + APIKey: "test-api-key", + TemplateID: "base", + }, server.URL) + if err != nil { + t.Fatalf("NewE2BSandbox failed: %v", err) + } + + ctx := context.Background() + if err := sb.Create(ctx); err != nil { + t.Fatalf("Create failed: %v", err) + } + + if sb.ID() != "sb-test-123" { + t.Errorf("expected sandbox ID 'sb-test-123', got %q", sb.ID()) + } +} + +func TestE2BSandbox_Exec(t *testing.T) { + var executedCmd string + + envdServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == "POST" && r.URL.Path == "/commands" { + var req map[string]interface{} + json.NewDecoder(r.Body).Decode(&req) + executedCmd = req["cmd"].(string) + + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(map[string]interface{}{ + "exitCode": 0, + "stdout": "hello world\n", + "stderr": "", + }) + return + } + w.WriteHeader(http.StatusNotFound) + })) + defer envdServer.Close() + + apiServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == "POST" && r.URL.Path == "/sandboxes" { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusCreated) + json.NewEncoder(w).Encode(map[string]interface{}{ + "sandboxID": "sb-test-123", + "templateID": "base", + "domain": "e2b.dev", + "envdAccessToken": "test-token", + }) + return + } + w.WriteHeader(http.StatusNotFound) + })) + defer apiServer.Close() + + sb, err := NewE2BSandboxWithAPIURL(Config{Mode: ModeE2B}, E2BConfig{ + APIKey: "test-key", + }, apiServer.URL) + if err != nil { + t.Fatalf("NewE2BSandbox failed: %v", err) + } + + sb.client.SetEnvdURL(envdServer.URL) + sb.created = true + sb.sandboxID = "sb-test-123" + + ctx := context.Background() + result, err := sb.Exec(ctx, "echo", "hello", "world") + if err != nil { + t.Fatalf("Exec failed: %v", err) + } + + if result.ExitCode != 0 { + t.Errorf("expected exit code 0, got %d", result.ExitCode) + } + + if result.Stdout != "hello world\n" { + t.Errorf("expected stdout 'hello world\\n', got %q", result.Stdout) + } + + if !strings.Contains(executedCmd, "echo hello world") { + t.Errorf("command not properly constructed: %q", executedCmd) + } +} + +func TestE2BSandbox_FileOperations(t *testing.T) { + fileStore := make(map[string][]byte) + + envdServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + path := r.URL.Query().Get("path") + + switch { + case r.Method == "GET" && r.URL.Path == "/files": + if content, ok := fileStore[path]; ok { + w.Write(content) + } else { + w.WriteHeader(http.StatusNotFound) + w.Write([]byte("file not found")) + } + + case r.Method == "POST" && r.URL.Path == "/files": + if err := r.ParseMultipartForm(10 << 20); err != nil { + w.WriteHeader(http.StatusBadRequest) + return + } + file, _, err := r.FormFile("file") + if err != nil { + w.WriteHeader(http.StatusBadRequest) + return + } + defer file.Close() + content, _ := io.ReadAll(file) + fileStore[path] = content + w.WriteHeader(http.StatusCreated) + + case r.Method == "DELETE" && r.URL.Path == "/files": + delete(fileStore, path) + w.WriteHeader(http.StatusNoContent) + + case r.Method == "POST" && r.URL.Path == "/commands": + var req map[string]interface{} + json.NewDecoder(r.Body).Decode(&req) + cmd := req["cmd"].(string) + + w.Header().Set("Content-Type", "application/json") + if strings.Contains(cmd, "test -e") { + parts := strings.Split(cmd, " ") + for i, p := range parts { + if p == "-e" && i+1 < len(parts) { + checkPath := strings.Split(parts[i+1], " ")[0] + if _, ok := fileStore[checkPath]; ok { + json.NewEncoder(w).Encode(map[string]interface{}{ + "exitCode": 0, + "stdout": "exists\n", + "stderr": "", + }) + return + } + } + } + json.NewEncoder(w).Encode(map[string]interface{}{ + "exitCode": 1, + "stdout": "", + "stderr": "", + }) + return + } + + json.NewEncoder(w).Encode(map[string]interface{}{ + "exitCode": 0, + "stdout": "", + "stderr": "", + }) + + default: + w.WriteHeader(http.StatusNotFound) + } + })) + defer envdServer.Close() + + apiServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusNotFound) + })) + defer apiServer.Close() + + sb, _ := NewE2BSandboxWithAPIURL(Config{Mode: ModeE2B}, E2BConfig{APIKey: "test-key"}, apiServer.URL) + sb.client.SetEnvdURL(envdServer.URL) + sb.created = true + sb.sandboxID = "sb-test-123" + + ctx := context.Background() + + content := []byte("test content") + if err := sb.WriteFile(ctx, "/test.txt", content, 0644); err != nil { + t.Fatalf("WriteFile failed: %v", err) + } + + readContent, err := sb.ReadFile(ctx, "/test.txt") + if err != nil { + t.Fatalf("ReadFile failed: %v", err) + } + if string(readContent) != string(content) { + t.Errorf("content mismatch: got %q, want %q", readContent, content) + } + + exists, err := sb.FileExists(ctx, "/test.txt") + if err != nil { + t.Fatalf("FileExists failed: %v", err) + } + if !exists { + t.Error("file should exist after write") + } + + if err := sb.DeleteFile(ctx, "/test.txt"); err != nil { + t.Fatalf("DeleteFile failed: %v", err) + } + + exists, err = sb.FileExists(ctx, "/test.txt") + if err != nil { + t.Fatalf("FileExists after delete failed: %v", err) + } + if exists { + t.Error("file should not exist after delete") + } +} + +func TestE2BSandbox_Destroy(t *testing.T) { + destroyed := false + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == "DELETE" && strings.HasPrefix(r.URL.Path, "/sandboxes/") { + destroyed = true + w.WriteHeader(http.StatusNoContent) + return + } + w.WriteHeader(http.StatusNotFound) + })) + defer server.Close() + + sb, _ := NewE2BSandboxWithAPIURL(Config{Mode: ModeE2B}, E2BConfig{APIKey: "test-key"}, server.URL) + sb.created = true + sb.sandboxID = "sb-test-123" + sb.client.Connect("sb-test-123", "test-token", "e2b.dev") + + ctx := context.Background() + if err := sb.Destroy(ctx); err != nil { + t.Fatalf("Destroy failed: %v", err) + } + + if !destroyed { + t.Error("sandbox should have been destroyed via API call") + } + + if sb.created { + t.Error("created flag should be false after destroy") + } + + if sb.sandboxID != "" { + t.Error("sandboxID should be empty after destroy") + } +} + +func TestE2BSandbox_RequiresAPIKey(t *testing.T) { + os.Unsetenv("E2B_API_KEY") + + _, err := NewE2BSandbox(Config{Mode: ModeE2B}, E2BConfig{}) + if err == nil { + t.Error("expected error when API key is missing") + } + + if !strings.Contains(err.Error(), "E2B_API_KEY") { + t.Errorf("error should mention E2B_API_KEY: %v", err) + } +} + +func TestE2BSandbox_Config(t *testing.T) { + os.Setenv("E2B_API_KEY", "test-key") + defer os.Unsetenv("E2B_API_KEY") + + cfg := Config{ + Mode: ModeE2B, + Image: "python", + Timeout: 10 * time.Minute, + WorkspacePath: "/workspace", + } + + sb, err := NewE2BSandbox(cfg, E2BConfig{}) + if err != nil { + t.Fatalf("NewE2BSandbox failed: %v", err) + } + + gotCfg := sb.Config() + if gotCfg.Mode != ModeE2B { + t.Errorf("expected mode %s, got %s", ModeE2B, gotCfg.Mode) + } + if gotCfg.Image != "python" { + t.Errorf("expected image 'python', got %q", gotCfg.Image) + } + if gotCfg.WorkspacePath != "/workspace" { + t.Errorf("expected workspace '/workspace', got %q", gotCfg.WorkspacePath) + } +} + +func TestFactory_CreateE2B(t *testing.T) { + os.Setenv("E2B_API_KEY", "test-key") + defer os.Unsetenv("E2B_API_KEY") + + factory := NewFactory(DefaultConfig()) + + sb, err := factory.Create(Config{Mode: ModeE2B}) + if err != nil { + t.Fatalf("Factory.Create failed: %v", err) + } + + _, ok := sb.(*E2BSandbox) + if !ok { + t.Error("expected E2BSandbox type") + } + + cfg := sb.Config() + if cfg.Mode != ModeE2B { + t.Errorf("expected mode %s, got %s", ModeE2B, cfg.Mode) + } +} diff --git a/pkg/harness/sandbox/host.go b/pkg/harness/sandbox/host.go new file mode 100644 index 00000000..ce4b2e41 --- /dev/null +++ b/pkg/harness/sandbox/host.go @@ -0,0 +1,385 @@ +package sandbox + +import ( + "bytes" + "context" + "fmt" + "io" + "os" + "os/exec" + "path/filepath" + "time" + + "github.com/google/uuid" +) + +type HostSandbox struct { + id string + config Config +} + +func NewHostSandbox(cfg Config) (*HostSandbox, error) { + return &HostSandbox{ + id: uuid.New().String(), + config: cfg, + }, nil +} + +func (s *HostSandbox) Create(ctx context.Context) error { + if s.config.WorkspacePath != "" { + return os.MkdirAll(s.config.WorkspacePath, 0755) + } + return nil +} + +func (s *HostSandbox) Exec(ctx context.Context, command string, args ...string) (*ExecResult, error) { + return s.ExecWithStdin(ctx, nil, command, args...) +} + +func (s *HostSandbox) ExecWithStdin(ctx context.Context, stdin io.Reader, command string, args ...string) (*ExecResult, error) { + start := time.Now() + + if s.config.Timeout > 0 { + var cancel context.CancelFunc + ctx, cancel = context.WithTimeout(ctx, s.config.Timeout) + defer cancel() + } + + cmd := exec.CommandContext(ctx, command, args...) + + if s.config.WorkspacePath != "" { + cmd.Dir = s.config.WorkspacePath + } + + for k, v := range s.config.Environment { + cmd.Env = append(cmd.Env, fmt.Sprintf("%s=%s", k, v)) + } + if len(cmd.Env) > 0 { + cmd.Env = append(os.Environ(), cmd.Env...) + } + + var stdout, stderr bytes.Buffer + cmd.Stdout = &stdout + cmd.Stderr = &stderr + + if stdin != nil { + cmd.Stdin = stdin + } + + err := cmd.Run() + duration := time.Since(start) + + result := &ExecResult{ + Stdout: stdout.String(), + Stderr: stderr.String(), + Duration: duration, + } + + if ctx.Err() == context.DeadlineExceeded { + result.Killed = true + result.KillReason = "timeout" + result.ExitCode = -1 + return result, nil + } + + if err != nil { + if exitErr, ok := err.(*exec.ExitError); ok { + result.ExitCode = exitErr.ExitCode() + } else { + return nil, err + } + } + + return result, nil +} + +func (s *HostSandbox) ReadFile(ctx context.Context, path string) ([]byte, error) { + fullPath := s.resolvePath(path) + return os.ReadFile(fullPath) +} + +func (s *HostSandbox) WriteFile(ctx context.Context, path string, content []byte, mode uint32) error { + fullPath := s.resolvePath(path) + + if err := os.MkdirAll(filepath.Dir(fullPath), 0755); err != nil { + return err + } + + if mode == 0 { + mode = 0644 + } + + return os.WriteFile(fullPath, content, os.FileMode(mode)) +} + +func (s *HostSandbox) DeleteFile(ctx context.Context, path string) error { + fullPath := s.resolvePath(path) + return os.Remove(fullPath) +} + +func (s *HostSandbox) ListFiles(ctx context.Context, path string) ([]FileInfo, error) { + fullPath := s.resolvePath(path) + + entries, err := os.ReadDir(fullPath) + if err != nil { + return nil, err + } + + var files []FileInfo + for _, entry := range entries { + info, err := entry.Info() + if err != nil { + continue + } + files = append(files, FileInfo{ + Name: entry.Name(), + Size: info.Size(), + Mode: uint32(info.Mode()), + ModTime: info.ModTime(), + IsDir: entry.IsDir(), + }) + } + + return files, nil +} + +func (s *HostSandbox) FileExists(ctx context.Context, path string) (bool, error) { + fullPath := s.resolvePath(path) + _, err := os.Stat(fullPath) + if err == nil { + return true, nil + } + if os.IsNotExist(err) { + return false, nil + } + return false, err +} + +func (s *HostSandbox) CopyIn(ctx context.Context, hostPath, sandboxPath string) error { + content, err := os.ReadFile(hostPath) + if err != nil { + return err + } + + info, err := os.Stat(hostPath) + if err != nil { + return err + } + + return s.WriteFile(ctx, sandboxPath, content, uint32(info.Mode())) +} + +func (s *HostSandbox) CopyOut(ctx context.Context, sandboxPath, hostPath string) error { + content, err := s.ReadFile(ctx, sandboxPath) + if err != nil { + return err + } + + if err := os.MkdirAll(filepath.Dir(hostPath), 0755); err != nil { + return err + } + + return os.WriteFile(hostPath, content, 0644) +} + +func (s *HostSandbox) GetMetrics(ctx context.Context) (*Metrics, error) { + return &Metrics{}, nil +} + +func (s *HostSandbox) Destroy(ctx context.Context) error { + return nil +} + +func (s *HostSandbox) ID() string { + return s.id +} + +func (s *HostSandbox) Config() *Config { + return &s.config +} + +func (s *HostSandbox) resolvePath(path string) string { + if filepath.IsAbs(path) { + return path + } + if s.config.WorkspacePath != "" { + return filepath.Join(s.config.WorkspacePath, path) + } + return path +} + +func (s *HostSandbox) ExecStream(ctx context.Context, opts ExecOptions, command string, args ...string) (ProcessHandle, error) { + timeout := opts.Timeout + if timeout == 0 { + timeout = s.config.Timeout + } + if timeout > 0 { + var cancel context.CancelFunc + ctx, cancel = context.WithTimeout(ctx, timeout) + _ = cancel + } + + cmd := exec.CommandContext(ctx, command, args...) + + cwd := opts.Cwd + if cwd == "" { + cwd = s.config.WorkspacePath + } + if cwd != "" { + cmd.Dir = cwd + } + + env := make(map[string]string) + for k, v := range s.config.Environment { + env[k] = v + } + for k, v := range opts.Env { + env[k] = v + } + for k, v := range env { + cmd.Env = append(cmd.Env, fmt.Sprintf("%s=%s", k, v)) + } + if len(cmd.Env) > 0 { + cmd.Env = append(os.Environ(), cmd.Env...) + } + + var stdinPipe io.WriteCloser + var err error + + if opts.Stdin != nil { + cmd.Stdin = opts.Stdin + } else { + stdinPipe, err = cmd.StdinPipe() + if err != nil { + return nil, fmt.Errorf("create stdin pipe: %w", err) + } + } + + stdoutPipe, err := cmd.StdoutPipe() + if err != nil { + return nil, fmt.Errorf("create stdout pipe: %w", err) + } + + stderrPipe, err := cmd.StderrPipe() + if err != nil { + return nil, fmt.Errorf("create stderr pipe: %w", err) + } + + if err := cmd.Start(); err != nil { + return nil, fmt.Errorf("start command: %w", err) + } + + handle := &hostProcessHandle{ + cmd: cmd, + ctx: ctx, + stdinPipe: stdinPipe, + start: time.Now(), + done: make(chan struct{}), + } + + go handle.streamOutput(stdoutPipe, opts.OnStdout, &handle.stdout) + go handle.streamOutput(stderrPipe, opts.OnStderr, &handle.stderr) + go handle.wait() + + return handle, nil +} + +func (s *HostSandbox) ListProcesses(ctx context.Context) ([]ProcessInfo, error) { + return nil, fmt.Errorf("ListProcesses not supported in host sandbox") +} + +func (s *HostSandbox) KillProcess(ctx context.Context, pid int) error { + proc, err := os.FindProcess(pid) + if err != nil { + return err + } + return proc.Kill() +} + +var _ StreamingSandbox = (*HostSandbox)(nil) + +type hostProcessHandle struct { + cmd *exec.Cmd + ctx context.Context + stdinPipe io.WriteCloser + start time.Time + done chan struct{} + result *ExecResult + err error + stdout bytes.Buffer + stderr bytes.Buffer +} + +func (h *hostProcessHandle) PID() int { + if h.cmd.Process != nil { + return h.cmd.Process.Pid + } + return 0 +} + +func (h *hostProcessHandle) Wait() (*ExecResult, error) { + <-h.done + if h.err != nil { + return h.result, h.err + } + return h.result, nil +} + +func (h *hostProcessHandle) Kill() error { + if h.cmd.Process != nil { + return h.cmd.Process.Kill() + } + return nil +} + +func (h *hostProcessHandle) SendStdin(data []byte) error { + if h.stdinPipe == nil { + return fmt.Errorf("stdin pipe not available") + } + _, err := h.stdinPipe.Write(data) + return err +} + +func (h *hostProcessHandle) streamOutput(r io.Reader, callback func([]byte), buf *bytes.Buffer) { + data := make([]byte, 4096) + for { + n, err := r.Read(data) + if n > 0 { + buf.Write(data[:n]) + if callback != nil { + callback(data[:n]) + } + } + if err != nil { + return + } + } +} + +func (h *hostProcessHandle) wait() { + defer close(h.done) + + err := h.cmd.Wait() + duration := time.Since(h.start) + + h.result = &ExecResult{ + Stdout: h.stdout.String(), + Stderr: h.stderr.String(), + Duration: duration, + } + + if h.ctx.Err() == context.DeadlineExceeded { + h.result.Killed = true + h.result.KillReason = "timeout" + h.result.ExitCode = -1 + return + } + + if err != nil { + if exitErr, ok := err.(*exec.ExitError); ok { + h.result.ExitCode = exitErr.ExitCode() + } else { + h.err = err + } + } +} diff --git a/pkg/harness/sandbox/sandbox.go b/pkg/harness/sandbox/sandbox.go new file mode 100644 index 00000000..4ee9c409 --- /dev/null +++ b/pkg/harness/sandbox/sandbox.go @@ -0,0 +1,224 @@ +package sandbox + +import ( + "context" + "fmt" + "io" + "time" +) + +type Mode string + +const ( + ModeHost Mode = "host" + ModeDocker Mode = "docker" + ModeE2B Mode = "e2b" +) + +type Config struct { + Mode Mode `json:"mode" yaml:"mode"` + Image string `json:"image,omitempty" yaml:"image,omitempty"` + Resources ResourceConfig `json:"resources,omitempty" yaml:"resources,omitempty"` + Network NetworkConfig `json:"network,omitempty" yaml:"network,omitempty"` + Filesystem FilesystemConfig `json:"filesystem,omitempty" yaml:"filesystem,omitempty"` + Timeout time.Duration `json:"timeout,omitempty" yaml:"timeout,omitempty"` + WorkspacePath string `json:"workspace_path,omitempty" yaml:"workspace_path,omitempty"` + Environment map[string]string `json:"environment,omitempty" yaml:"environment,omitempty"` + // RegistryAuth holds authentication for private container registries (ECR, GCR, Docker Hub, etc.) + RegistryAuth *RegistryAuthConfig `json:"registry_auth,omitempty" yaml:"registry_auth,omitempty"` +} + +// RegistryAuthConfig holds authentication credentials for private container registries. +// Supports multiple authentication methods: +// - Username/Password: Basic auth for Docker Hub, self-hosted registries +// - IdentityToken: OAuth tokens for cloud providers (ECR, GCR, ACR) +// - DockerConfigPath: Path to ~/.docker/config.json for complex setups +type RegistryAuthConfig struct { + // Username for basic auth (Docker Hub, self-hosted registries) + Username string `json:"username,omitempty" yaml:"username,omitempty"` + // Password or access token for basic auth + Password string `json:"password,omitempty" yaml:"password,omitempty"` + // IdentityToken is an OAuth bearer token (for cloud providers like ECR, GCR) + IdentityToken string `json:"identity_token,omitempty" yaml:"identity_token,omitempty"` + // ServerAddress is the registry server (e.g., "https://index.docker.io/v1/", "123456789.dkr.ecr.us-east-1.amazonaws.com") + ServerAddress string `json:"server_address,omitempty" yaml:"server_address,omitempty"` + // DockerConfigPath is the path to a Docker config.json file for complex auth setups + // When set, this takes precedence over Username/Password/IdentityToken + DockerConfigPath string `json:"docker_config_path,omitempty" yaml:"docker_config_path,omitempty"` +} + +type ResourceConfig struct { + CPU float64 `json:"cpu,omitempty" yaml:"cpu,omitempty"` + Memory string `json:"memory,omitempty" yaml:"memory,omitempty"` + Disk string `json:"disk,omitempty" yaml:"disk,omitempty"` + PIDs int64 `json:"pids,omitempty" yaml:"pids,omitempty"` +} + +type NetworkConfig struct { + Enabled bool `json:"enabled" yaml:"enabled"` + AllowedHosts []string `json:"allowed_hosts,omitempty" yaml:"allowed_hosts,omitempty"` + AllowedPorts []int `json:"allowed_ports,omitempty" yaml:"allowed_ports,omitempty"` + DNSServers []string `json:"dns_servers,omitempty" yaml:"dns_servers,omitempty"` +} + +type FilesystemConfig struct { + ReadOnly []string `json:"read_only,omitempty" yaml:"read_only,omitempty"` + ReadWrite []string `json:"read_write,omitempty" yaml:"read_write,omitempty"` + Denied []string `json:"denied,omitempty" yaml:"denied,omitempty"` + TempDir string `json:"temp_dir,omitempty" yaml:"temp_dir,omitempty"` +} + +type ExecResult struct { + ExitCode int `json:"exit_code"` + Stdout string `json:"stdout"` + Stderr string `json:"stderr"` + Duration time.Duration `json:"duration"` + Killed bool `json:"killed,omitempty"` + KillReason string `json:"kill_reason,omitempty"` +} + +type Metrics struct { + CPUUsage float64 `json:"cpu_usage"` + MemoryUsageBytes int64 `json:"memory_usage_bytes"` + MemoryMaxBytes int64 `json:"memory_max_bytes"` + DiskReadBytes int64 `json:"disk_read_bytes"` + DiskWriteBytes int64 `json:"disk_write_bytes"` + NetworkRxBytes int64 `json:"network_rx_bytes"` + NetworkTxBytes int64 `json:"network_tx_bytes"` + ProcessCount int `json:"process_count"` +} + +type Sandbox interface { + Create(ctx context.Context) error + Exec(ctx context.Context, command string, args ...string) (*ExecResult, error) + ExecWithStdin(ctx context.Context, stdin io.Reader, command string, args ...string) (*ExecResult, error) + ReadFile(ctx context.Context, path string) ([]byte, error) + WriteFile(ctx context.Context, path string, content []byte, mode uint32) error + DeleteFile(ctx context.Context, path string) error + ListFiles(ctx context.Context, path string) ([]FileInfo, error) + FileExists(ctx context.Context, path string) (bool, error) + CopyIn(ctx context.Context, hostPath, sandboxPath string) error + CopyOut(ctx context.Context, sandboxPath, hostPath string) error + GetMetrics(ctx context.Context) (*Metrics, error) + Destroy(ctx context.Context) error + ID() string + Config() *Config +} + +type FileInfo struct { + Name string `json:"name"` + Size int64 `json:"size"` + Mode uint32 `json:"mode"` + ModTime time.Time `json:"mod_time"` + IsDir bool `json:"is_dir"` +} + +type ExecOptions struct { + Cwd string + Env map[string]string + Timeout time.Duration + OnStdout func(data []byte) + OnStderr func(data []byte) + Stdin io.Reader +} + +type ProcessInfo struct { + PID int `json:"pid"` + Command string `json:"command"` + Args []string `json:"args"` + Cwd string `json:"cwd"` + Env map[string]string `json:"env"` +} + +type ProcessHandle interface { + PID() int + Wait() (*ExecResult, error) + Kill() error + SendStdin(data []byte) error +} + +type StreamingSandbox interface { + Sandbox + ExecStream(ctx context.Context, opts ExecOptions, command string, args ...string) (ProcessHandle, error) + ListProcesses(ctx context.Context) ([]ProcessInfo, error) + KillProcess(ctx context.Context, pid int) error +} + +type Factory struct { + DefaultConfig Config +} + +func NewFactory(defaults Config) *Factory { + return &Factory{DefaultConfig: defaults} +} + +func (f *Factory) Create(cfg Config) (Sandbox, error) { + merged := f.mergeConfig(cfg) + + switch merged.Mode { + case ModeHost, "": + return NewHostSandbox(merged) + case ModeDocker: + return NewDockerSandbox(merged) + case ModeE2B: + return NewE2BSandbox(merged, E2BConfig{ + TemplateID: merged.Image, + TimeoutSec: int(merged.Timeout.Seconds()), + }) + default: + return nil, fmt.Errorf("unknown sandbox mode: %s", merged.Mode) + } +} + +func (f *Factory) mergeConfig(cfg Config) Config { + merged := cfg + + if merged.Mode == "" { + merged.Mode = f.DefaultConfig.Mode + if merged.Mode == "" { + merged.Mode = ModeHost + } + } + + if merged.Timeout == 0 { + merged.Timeout = f.DefaultConfig.Timeout + if merged.Timeout == 0 { + merged.Timeout = 30 * time.Minute + } + } + + if merged.Image == "" { + merged.Image = f.DefaultConfig.Image + } + + if merged.Resources.CPU == 0 { + merged.Resources.CPU = f.DefaultConfig.Resources.CPU + } + if merged.Resources.Memory == "" { + merged.Resources.Memory = f.DefaultConfig.Resources.Memory + } + + if merged.WorkspacePath == "" { + merged.WorkspacePath = f.DefaultConfig.WorkspacePath + } + + return merged +} + +func DefaultConfig() Config { + return Config{ + Mode: ModeHost, + Timeout: 30 * time.Minute, + Resources: ResourceConfig{ + CPU: 2, + Memory: "4g", + PIDs: 1000, + }, + Network: NetworkConfig{ + Enabled: false, + }, + Filesystem: FilesystemConfig{ + TempDir: "/tmp", + }, + } +} diff --git a/pkg/harness/sandbox/sandbox_test.go b/pkg/harness/sandbox/sandbox_test.go new file mode 100644 index 00000000..1e67e23e --- /dev/null +++ b/pkg/harness/sandbox/sandbox_test.go @@ -0,0 +1,347 @@ +package sandbox + +import ( + "context" + "os" + "path/filepath" + "strings" + "testing" + "time" +) + +func TestHostSandbox_Create(t *testing.T) { + tmpDir := t.TempDir() + + cfg := Config{ + Mode: ModeHost, + WorkspacePath: filepath.Join(tmpDir, "workspace"), + } + + sb, err := NewHostSandbox(cfg) + if err != nil { + t.Fatalf("NewHostSandbox failed: %v", err) + } + + ctx := context.Background() + if err := sb.Create(ctx); err != nil { + t.Fatalf("Create failed: %v", err) + } + + if _, err := os.Stat(cfg.WorkspacePath); os.IsNotExist(err) { + t.Error("workspace directory was not created") + } +} + +func TestHostSandbox_Exec(t *testing.T) { + sb, _ := NewHostSandbox(Config{Mode: ModeHost}) + + ctx := context.Background() + result, err := sb.Exec(ctx, "echo", "hello") + if err != nil { + t.Fatalf("Exec failed: %v", err) + } + + if result.ExitCode != 0 { + t.Errorf("expected exit code 0, got %d", result.ExitCode) + } + + if result.Stdout != "hello\n" { + t.Errorf("expected stdout 'hello\\n', got %q", result.Stdout) + } +} + +func TestHostSandbox_ExecTimeout(t *testing.T) { + sb, _ := NewHostSandbox(Config{ + Mode: ModeHost, + Timeout: 100 * time.Millisecond, + }) + + ctx := context.Background() + result, err := sb.Exec(ctx, "sleep", "10") + if err != nil { + t.Fatalf("Exec failed: %v", err) + } + + if !result.Killed { + t.Error("expected process to be killed") + } + + if result.KillReason != "timeout" { + t.Errorf("expected kill reason 'timeout', got %q", result.KillReason) + } +} + +func TestHostSandbox_FileOperations(t *testing.T) { + tmpDir := t.TempDir() + + sb, _ := NewHostSandbox(Config{ + Mode: ModeHost, + WorkspacePath: tmpDir, + }) + + ctx := context.Background() + + content := []byte("test content") + if err := sb.WriteFile(ctx, "test.txt", content, 0644); err != nil { + t.Fatalf("WriteFile failed: %v", err) + } + + exists, err := sb.FileExists(ctx, "test.txt") + if err != nil { + t.Fatalf("FileExists failed: %v", err) + } + if !exists { + t.Error("file should exist after write") + } + + readContent, err := sb.ReadFile(ctx, "test.txt") + if err != nil { + t.Fatalf("ReadFile failed: %v", err) + } + if string(readContent) != string(content) { + t.Errorf("content mismatch: got %q, want %q", readContent, content) + } + + files, err := sb.ListFiles(ctx, ".") + if err != nil { + t.Fatalf("ListFiles failed: %v", err) + } + + found := false + for _, f := range files { + if f.Name == "test.txt" { + found = true + break + } + } + if !found { + t.Error("test.txt not found in file listing") + } + + if err := sb.DeleteFile(ctx, "test.txt"); err != nil { + t.Fatalf("DeleteFile failed: %v", err) + } + + exists, _ = sb.FileExists(ctx, "test.txt") + if exists { + t.Error("file should not exist after delete") + } +} + +func TestHostSandbox_CopyInOut(t *testing.T) { + tmpDir := t.TempDir() + workspaceDir := filepath.Join(tmpDir, "workspace") + + sb, _ := NewHostSandbox(Config{ + Mode: ModeHost, + WorkspacePath: workspaceDir, + }) + + ctx := context.Background() + sb.Create(ctx) + + hostFile := filepath.Join(tmpDir, "source.txt") + content := []byte("copy test") + os.WriteFile(hostFile, content, 0644) + + if err := sb.CopyIn(ctx, hostFile, "copied.txt"); err != nil { + t.Fatalf("CopyIn failed: %v", err) + } + + exists, _ := sb.FileExists(ctx, "copied.txt") + if !exists { + t.Error("copied file should exist in sandbox") + } + + destFile := filepath.Join(tmpDir, "dest.txt") + if err := sb.CopyOut(ctx, "copied.txt", destFile); err != nil { + t.Fatalf("CopyOut failed: %v", err) + } + + destContent, err := os.ReadFile(destFile) + if err != nil { + t.Fatalf("failed to read dest file: %v", err) + } + if string(destContent) != string(content) { + t.Errorf("content mismatch after copy out") + } +} + +func TestHostSandbox_Environment(t *testing.T) { + sb, _ := NewHostSandbox(Config{ + Mode: ModeHost, + Environment: map[string]string{ + "TEST_VAR": "test_value", + }, + }) + + ctx := context.Background() + result, err := sb.Exec(ctx, "sh", "-c", "echo $TEST_VAR") + if err != nil { + t.Fatalf("Exec failed: %v", err) + } + + if result.Stdout != "test_value\n" { + t.Errorf("expected 'test_value\\n', got %q", result.Stdout) + } +} + +func TestFactory_CreateHost(t *testing.T) { + factory := NewFactory(DefaultConfig()) + + sb, err := factory.Create(Config{Mode: ModeHost}) + if err != nil { + t.Fatalf("Factory.Create failed: %v", err) + } + + if sb.ID() == "" { + t.Error("sandbox should have an ID") + } + + cfg := sb.Config() + if cfg.Mode != ModeHost { + t.Errorf("expected mode %s, got %s", ModeHost, cfg.Mode) + } +} + +func TestFactory_DefaultMode(t *testing.T) { + factory := NewFactory(DefaultConfig()) + + sb, err := factory.Create(Config{}) + if err != nil { + t.Fatalf("Factory.Create failed: %v", err) + } + + cfg := sb.Config() + if cfg.Mode != ModeHost { + t.Errorf("expected default mode %s, got %s", ModeHost, cfg.Mode) + } +} + +func TestFactory_MergeConfig(t *testing.T) { + factory := NewFactory(Config{ + Mode: ModeHost, + Timeout: 5 * time.Minute, + Resources: ResourceConfig{ + CPU: 4, + Memory: "8Gi", + }, + }) + + sb, _ := factory.Create(Config{ + Resources: ResourceConfig{ + CPU: 2, + }, + }) + + cfg := sb.Config() + + if cfg.Timeout != 5*time.Minute { + t.Errorf("timeout should inherit default: got %v", cfg.Timeout) + } + + if cfg.Resources.CPU != 2 { + t.Errorf("CPU should be overridden: got %v", cfg.Resources.CPU) + } + + if cfg.Resources.Memory != "8Gi" { + t.Errorf("Memory should inherit default: got %v", cfg.Resources.Memory) + } +} + +func TestDefaultConfig(t *testing.T) { + cfg := DefaultConfig() + + if cfg.Mode != ModeHost { + t.Errorf("expected default mode %s", ModeHost) + } + + if cfg.Timeout != 30*time.Minute { + t.Errorf("expected default timeout 30m, got %v", cfg.Timeout) + } + + if cfg.Resources.CPU != 2 { + t.Errorf("expected default CPU 2, got %v", cfg.Resources.CPU) + } + + if cfg.Network.Enabled { + t.Error("network should be disabled by default") + } +} + +func TestHostSandbox_ExecStream(t *testing.T) { + sb, _ := NewHostSandbox(Config{Mode: ModeHost}) + + ctx := context.Background() + var stdoutData []byte + + handle, err := sb.ExecStream(ctx, ExecOptions{ + OnStdout: func(data []byte) { + stdoutData = append(stdoutData, data...) + }, + }, "echo", "streaming test") + if err != nil { + t.Fatalf("ExecStream failed: %v", err) + } + + result, err := handle.Wait() + if err != nil { + t.Fatalf("Wait failed: %v", err) + } + + if result.ExitCode != 0 { + t.Errorf("expected exit code 0, got %d", result.ExitCode) + } + + if !strings.Contains(string(stdoutData), "streaming test") { + t.Errorf("expected stdout callback to receive 'streaming test', got %q", string(stdoutData)) + } +} + +func TestHostSandbox_ExecStream_WithStderr(t *testing.T) { + sb, _ := NewHostSandbox(Config{Mode: ModeHost}) + + ctx := context.Background() + var stderrData []byte + + handle, err := sb.ExecStream(ctx, ExecOptions{ + OnStderr: func(data []byte) { + stderrData = append(stderrData, data...) + }, + }, "sh", "-c", "echo error >&2") + if err != nil { + t.Fatalf("ExecStream failed: %v", err) + } + + result, err := handle.Wait() + if err != nil { + t.Fatalf("Wait failed: %v", err) + } + + if result.ExitCode != 0 { + t.Errorf("expected exit code 0, got %d", result.ExitCode) + } + + if !strings.Contains(string(stderrData), "error") { + t.Errorf("expected stderr callback to receive 'error', got %q", string(stderrData)) + } +} + +func TestHostSandbox_KillProcess(t *testing.T) { + sb, _ := NewHostSandbox(Config{Mode: ModeHost}) + + ctx := context.Background() + handle, err := sb.ExecStream(ctx, ExecOptions{}, "sleep", "60") + if err != nil { + t.Fatalf("ExecStream failed: %v", err) + } + + if err := handle.Kill(); err != nil { + t.Fatalf("Kill failed: %v", err) + } + + result, _ := handle.Wait() + if !result.Killed && result.ExitCode >= 0 { + t.Error("expected process to be killed") + } +} diff --git a/pkg/harness/session/history_store.go b/pkg/harness/session/history_store.go new file mode 100644 index 00000000..e9d3c14b --- /dev/null +++ b/pkg/harness/session/history_store.go @@ -0,0 +1,176 @@ +package session + +import ( + "encoding/json" + "fmt" + "os" + "path/filepath" + "sync" + "time" + + "github.com/firebase/genkit/go/ai" +) + +// HistoryStore persists message history for sessions. +// This enables REPL-style prolonged conversations with the agent. +type HistoryStore struct { + basePath string + mu sync.RWMutex +} + +// StoredMessage is a serializable version of ai.Message +type StoredMessage struct { + Role string `json:"role"` + Content string `json:"content"` + ToolCalls []StoredToolCall `json:"tool_calls,omitempty"` + Metadata map[string]interface{} `json:"metadata,omitempty"` + Timestamp time.Time `json:"timestamp"` +} + +// StoredToolCall is a serializable version of tool call data +type StoredToolCall struct { + Name string `json:"name"` + Arguments map[string]interface{} `json:"arguments"` + Result string `json:"result,omitempty"` +} + +// SessionHistory contains all messages for a session +type SessionHistory struct { + SessionID string `json:"session_id"` + Messages []StoredMessage `json:"messages"` + TotalTokens int `json:"total_tokens"` + CreatedAt time.Time `json:"created_at"` + UpdatedAt time.Time `json:"updated_at"` + Metadata map[string]interface{} `json:"metadata,omitempty"` +} + +// NewHistoryStore creates a new history store +func NewHistoryStore(basePath string) *HistoryStore { + return &HistoryStore{ + basePath: basePath, + } +} + +// historyPath returns the path for a session's history file +func (h *HistoryStore) historyPath(sessionID string) string { + return filepath.Join(h.basePath, "session", sessionID, ".history.json") +} + +// Load loads message history for a session +func (h *HistoryStore) Load(sessionID string) (*SessionHistory, error) { + h.mu.RLock() + defer h.mu.RUnlock() + + path := h.historyPath(sessionID) + data, err := os.ReadFile(path) + if err != nil { + if os.IsNotExist(err) { + // Return empty history for new sessions + return &SessionHistory{ + SessionID: sessionID, + Messages: []StoredMessage{}, + CreatedAt: time.Now(), + UpdatedAt: time.Now(), + }, nil + } + return nil, fmt.Errorf("failed to read history: %w", err) + } + + var history SessionHistory + if err := json.Unmarshal(data, &history); err != nil { + return nil, fmt.Errorf("failed to parse history: %w", err) + } + + return &history, nil +} + +// Save persists message history for a session +func (h *HistoryStore) Save(sessionID string, history *SessionHistory) error { + h.mu.Lock() + defer h.mu.Unlock() + + history.UpdatedAt = time.Now() + + data, err := json.MarshalIndent(history, "", " ") + if err != nil { + return fmt.Errorf("failed to marshal history: %w", err) + } + + path := h.historyPath(sessionID) + dir := filepath.Dir(path) + if err := os.MkdirAll(dir, 0755); err != nil { + return fmt.Errorf("failed to create history directory: %w", err) + } + + if err := os.WriteFile(path, data, 0644); err != nil { + return fmt.Errorf("failed to write history: %w", err) + } + + return nil +} + +// Append adds messages to a session's history +func (h *HistoryStore) Append(sessionID string, messages []StoredMessage) error { + history, err := h.Load(sessionID) + if err != nil { + return err + } + + history.Messages = append(history.Messages, messages...) + return h.Save(sessionID, history) +} + +// Clear removes all messages from a session's history +func (h *HistoryStore) Clear(sessionID string) error { + h.mu.Lock() + defer h.mu.Unlock() + + return os.Remove(h.historyPath(sessionID)) +} + +// ToAIMessages converts stored messages back to ai.Message format +func (h *SessionHistory) ToAIMessages() []*ai.Message { + messages := make([]*ai.Message, len(h.Messages)) + for i, msg := range h.Messages { + var parts []*ai.Part + if msg.Content != "" { + parts = append(parts, ai.NewTextPart(msg.Content)) + } + + // Note: Tool calls would need to be reconstructed from StoredToolCall + // This is a simplified conversion + + messages[i] = &ai.Message{ + Role: ai.Role(msg.Role), + Content: parts, + } + } + return messages +} + +// FromAIMessage converts an ai.Message to StoredMessage +func FromAIMessage(msg *ai.Message) StoredMessage { + stored := StoredMessage{ + Role: string(msg.Role), + Timestamp: time.Now(), + } + + // Extract text content + for _, part := range msg.Content { + if part.IsText() { + stored.Content += part.Text + } + // TODO: Handle tool request parts + } + + return stored +} + +// FromAIMessages converts multiple ai.Messages to StoredMessages +func FromAIMessages(msgs []*ai.Message) []StoredMessage { + stored := make([]StoredMessage, len(msgs)) + for i, msg := range msgs { + stored[i] = FromAIMessage(msg) + } + return stored +} diff --git a/pkg/harness/session/manager.go b/pkg/harness/session/manager.go new file mode 100644 index 00000000..ecf7eb49 --- /dev/null +++ b/pkg/harness/session/manager.go @@ -0,0 +1,363 @@ +package session + +import ( + "context" + "encoding/json" + "fmt" + "os" + "path/filepath" + "syscall" + "time" +) + +type Manager struct { + basePath string + lockTimeout time.Duration +} + +type Session struct { + ID string `json:"id"` + Path string `json:"path"` + RepoURL string `json:"repo_url,omitempty"` + Branch string `json:"branch,omitempty"` + CreatedAt time.Time `json:"created_at"` + LastUsedAt time.Time `json:"last_used_at"` + TotalRuns int `json:"total_runs"` + IsLocked bool `json:"is_locked"` + LockedBy string `json:"locked_by,omitempty"` + LockedAt time.Time `json:"locked_at,omitempty"` +} + +type LockInfo struct { + RunID string `json:"run_id"` + PID int `json:"pid"` + Hostname string `json:"hostname"` + AcquiredAt time.Time `json:"acquired_at"` + ExpiresAt time.Time `json:"expires_at"` +} + +type SessionMeta struct { + ID string `json:"id"` + RepoURL string `json:"repo_url,omitempty"` + Branch string `json:"branch,omitempty"` + CreatedAt time.Time `json:"created_at"` + LastUsedAt time.Time `json:"last_used_at"` + TotalRuns int `json:"total_runs"` +} + +func NewManager(basePath string) *Manager { + return &Manager{ + basePath: basePath, + lockTimeout: 1 * time.Hour, + } +} + +func (m *Manager) SetLockTimeout(d time.Duration) { + m.lockTimeout = d +} + +func (m *Manager) SessionPath(sessionID string) string { + return filepath.Join(m.basePath, "session", sessionID) +} + +func (m *Manager) lockPath(sessionID string) string { + return filepath.Join(m.SessionPath(sessionID), ".session.lock") +} + +func (m *Manager) metaPath(sessionID string) string { + return filepath.Join(m.SessionPath(sessionID), ".session.meta") +} + +func (m *Manager) Create(ctx context.Context, sessionID string, repoURL string, branch string) (*Session, error) { + sessionPath := m.SessionPath(sessionID) + + if err := os.MkdirAll(sessionPath, 0755); err != nil { + return nil, fmt.Errorf("failed to create session directory: %w", err) + } + + meta := SessionMeta{ + ID: sessionID, + RepoURL: repoURL, + Branch: branch, + CreatedAt: time.Now(), + LastUsedAt: time.Now(), + TotalRuns: 0, + } + + if err := m.writeMeta(sessionID, &meta); err != nil { + return nil, err + } + + return &Session{ + ID: sessionID, + Path: sessionPath, + RepoURL: repoURL, + Branch: branch, + CreatedAt: meta.CreatedAt, + LastUsedAt: meta.LastUsedAt, + TotalRuns: 0, + IsLocked: false, + }, nil +} + +func (m *Manager) Get(ctx context.Context, sessionID string) (*Session, error) { + sessionPath := m.SessionPath(sessionID) + + if _, err := os.Stat(sessionPath); os.IsNotExist(err) { + return nil, fmt.Errorf("session not found: %s", sessionID) + } + + meta, err := m.readMeta(sessionID) + if err != nil { + meta = &SessionMeta{ + ID: sessionID, + CreatedAt: time.Now(), + } + } + + session := &Session{ + ID: sessionID, + Path: sessionPath, + RepoURL: meta.RepoURL, + Branch: meta.Branch, + CreatedAt: meta.CreatedAt, + LastUsedAt: meta.LastUsedAt, + TotalRuns: meta.TotalRuns, + } + + lockInfo, err := m.getLockInfo(sessionID) + if err == nil && lockInfo != nil { + if time.Now().Before(lockInfo.ExpiresAt) && m.isProcessAlive(lockInfo.PID) { + session.IsLocked = true + session.LockedBy = lockInfo.RunID + session.LockedAt = lockInfo.AcquiredAt + } else { + _ = os.Remove(m.lockPath(sessionID)) + } + } + + return session, nil +} + +func (m *Manager) List(ctx context.Context) ([]*Session, error) { + sessionDir := filepath.Join(m.basePath, "session") + + if _, err := os.Stat(sessionDir); os.IsNotExist(err) { + return []*Session{}, nil + } + + entries, err := os.ReadDir(sessionDir) + if err != nil { + return nil, fmt.Errorf("failed to read sessions directory: %w", err) + } + + var sessions []*Session + for _, entry := range entries { + if !entry.IsDir() { + continue + } + + session, err := m.Get(ctx, entry.Name()) + if err != nil { + continue + } + sessions = append(sessions, session) + } + + return sessions, nil +} + +func (m *Manager) Delete(ctx context.Context, sessionID string) error { + session, err := m.Get(ctx, sessionID) + if err != nil { + return err + } + + if session.IsLocked { + return fmt.Errorf("cannot delete locked session (locked by %s)", session.LockedBy) + } + + return os.RemoveAll(session.Path) +} + +func (m *Manager) Cleanup(ctx context.Context, olderThan time.Duration) ([]string, error) { + sessions, err := m.List(ctx) + if err != nil { + return nil, err + } + + cutoff := time.Now().Add(-olderThan) + var deleted []string + + for _, session := range sessions { + if session.LastUsedAt.Before(cutoff) && !session.IsLocked { + if err := os.RemoveAll(session.Path); err == nil { + deleted = append(deleted, session.ID) + } + } + } + + return deleted, nil +} + +func (m *Manager) AcquireLock(ctx context.Context, sessionID string, runID string) error { + session, err := m.Get(ctx, sessionID) + if err != nil { + return err + } + + if session.IsLocked { + return fmt.Errorf("session already locked by run %s", session.LockedBy) + } + + hostname, _ := os.Hostname() + lockInfo := LockInfo{ + RunID: runID, + PID: os.Getpid(), + Hostname: hostname, + AcquiredAt: time.Now(), + ExpiresAt: time.Now().Add(m.lockTimeout), + } + + data, err := json.MarshalIndent(lockInfo, "", " ") + if err != nil { + return fmt.Errorf("failed to marshal lock info: %w", err) + } + + lockPath := m.lockPath(sessionID) + f, err := os.OpenFile(lockPath, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0644) + if err != nil { + if os.IsExist(err) { + return fmt.Errorf("session lock already exists") + } + return fmt.Errorf("failed to create lock file: %w", err) + } + defer f.Close() + + if _, err := f.Write(data); err != nil { + os.Remove(lockPath) + return fmt.Errorf("failed to write lock file: %w", err) + } + + meta, _ := m.readMeta(sessionID) + if meta != nil { + meta.LastUsedAt = time.Now() + meta.TotalRuns++ + _ = m.writeMeta(sessionID, meta) + } + + return nil +} + +func (m *Manager) ReleaseLock(ctx context.Context, sessionID string, runID string) error { + lockInfo, err := m.getLockInfo(sessionID) + if err != nil { + return nil + } + + if lockInfo.RunID != runID { + return fmt.Errorf("lock held by different run: %s", lockInfo.RunID) + } + + return os.Remove(m.lockPath(sessionID)) +} + +func (m *Manager) ForceUnlock(ctx context.Context, sessionID string) error { + return os.Remove(m.lockPath(sessionID)) +} + +func (m *Manager) RefreshLock(ctx context.Context, sessionID string, runID string) error { + lockInfo, err := m.getLockInfo(sessionID) + if err != nil { + return fmt.Errorf("no lock to refresh: %w", err) + } + + if lockInfo.RunID != runID { + return fmt.Errorf("lock held by different run: %s", lockInfo.RunID) + } + + lockInfo.ExpiresAt = time.Now().Add(m.lockTimeout) + + data, err := json.MarshalIndent(lockInfo, "", " ") + if err != nil { + return err + } + + return os.WriteFile(m.lockPath(sessionID), data, 0644) +} + +func (m *Manager) GetOrCreate(ctx context.Context, sessionID string, repoURL string, branch string) (*Session, error) { + session, err := m.Get(ctx, sessionID) + if err == nil { + return session, nil + } + + return m.Create(ctx, sessionID, repoURL, branch) +} + +func (m *Manager) getLockInfo(sessionID string) (*LockInfo, error) { + data, err := os.ReadFile(m.lockPath(sessionID)) + if err != nil { + return nil, err + } + + var lockInfo LockInfo + if err := json.Unmarshal(data, &lockInfo); err != nil { + return nil, err + } + + return &lockInfo, nil +} + +func (m *Manager) readMeta(sessionID string) (*SessionMeta, error) { + data, err := os.ReadFile(m.metaPath(sessionID)) + if err != nil { + return nil, err + } + + var meta SessionMeta + if err := json.Unmarshal(data, &meta); err != nil { + return nil, err + } + + return &meta, nil +} + +func (m *Manager) writeMeta(sessionID string, meta *SessionMeta) error { + data, err := json.MarshalIndent(meta, "", " ") + if err != nil { + return err + } + + return os.WriteFile(m.metaPath(sessionID), data, 0644) +} + +func (m *Manager) isProcessAlive(pid int) bool { + if pid <= 0 { + return false + } + + process, err := os.FindProcess(pid) + if err != nil { + return false + } + + err = process.Signal(syscall.Signal(0)) + return err == nil +} + +func (m *Manager) DiskUsage(sessionID string) (int64, error) { + var size int64 + + err := filepath.Walk(m.SessionPath(sessionID), func(_ string, info os.FileInfo, err error) error { + if err != nil { + return err + } + if !info.IsDir() { + size += info.Size() + } + return nil + }) + + return size, err +} diff --git a/pkg/harness/session/session_executor.go b/pkg/harness/session/session_executor.go new file mode 100644 index 00000000..170b211b --- /dev/null +++ b/pkg/harness/session/session_executor.go @@ -0,0 +1,255 @@ +package session + +import ( + "context" + "fmt" + "time" + + "station/pkg/harness" + "station/pkg/harness/sandbox" + "station/pkg/harness/workspace" + + "github.com/firebase/genkit/go/ai" + "github.com/firebase/genkit/go/genkit" +) + +// SessionExecutor wraps AgenticExecutor with session persistence. +// It maintains message history and workspace state across multiple interactions. +type SessionExecutor struct { + genkitApp *genkit.Genkit + sessionManager *Manager + historyStore *HistoryStore + harnessConfig *harness.HarnessConfig + agentConfig *harness.AgentHarnessConfig + modelName string + systemPrompt string + + // Active session state + currentSessionID string + currentSession *Session + currentWorkspace harness.WorkspaceManager + currentSandbox sandbox.Sandbox + sandboxConfig *sandbox.Config + tools []ai.Tool +} + +// SessionExecutorConfig holds configuration for the session executor +type SessionExecutorConfig struct { + BasePath string + HarnessConfig *harness.HarnessConfig + AgentConfig *harness.AgentHarnessConfig + ModelName string + SystemPrompt string + SandboxConfig *sandbox.Config +} + +// NewSessionExecutor creates a new session executor +func NewSessionExecutor(genkitApp *genkit.Genkit, cfg SessionExecutorConfig) *SessionExecutor { + if cfg.HarnessConfig == nil { + cfg.HarnessConfig = harness.DefaultHarnessConfig() + } + if cfg.AgentConfig == nil { + cfg.AgentConfig = harness.DefaultAgentHarnessConfig() + } + if cfg.ModelName == "" { + cfg.ModelName = "openai/gpt-4o-mini" + } + + return &SessionExecutor{ + genkitApp: genkitApp, + sessionManager: NewManager(cfg.BasePath), + historyStore: NewHistoryStore(cfg.BasePath), + harnessConfig: cfg.HarnessConfig, + agentConfig: cfg.AgentConfig, + modelName: cfg.ModelName, + systemPrompt: cfg.SystemPrompt, + sandboxConfig: cfg.SandboxConfig, + } +} + +// StartSession starts or resumes a session +func (se *SessionExecutor) StartSession(ctx context.Context, sessionID string, repoURL string, branch string) error { + // Get or create session + session, err := se.sessionManager.GetOrCreate(ctx, sessionID, repoURL, branch) + if err != nil { + return fmt.Errorf("failed to get/create session: %w", err) + } + + // Acquire lock + runID := fmt.Sprintf("repl-%d", time.Now().UnixNano()) + if err := se.sessionManager.AcquireLock(ctx, sessionID, runID); err != nil { + return fmt.Errorf("failed to acquire session lock: %w", err) + } + + // Initialize workspace + ws := workspace.NewHostWorkspace(session.Path) + if err := ws.Initialize(ctx); err != nil { + se.sessionManager.ReleaseLock(ctx, sessionID, runID) + return fmt.Errorf("failed to initialize workspace: %w", err) + } + + // Initialize sandbox if configured (keep alive for session duration) + var sb sandbox.Sandbox + if se.sandboxConfig != nil { + // Set workspace path for sandbox + sandboxCfg := *se.sandboxConfig + sandboxCfg.WorkspacePath = session.Path + + factory := sandbox.NewFactory(sandbox.DefaultConfig()) + sb, err = factory.Create(sandboxCfg) + if err != nil { + se.sessionManager.ReleaseLock(ctx, sessionID, runID) + return fmt.Errorf("failed to create sandbox: %w", err) + } + + if err := sb.Create(ctx); err != nil { + se.sessionManager.ReleaseLock(ctx, sessionID, runID) + return fmt.Errorf("failed to initialize sandbox: %w", err) + } + } + + se.currentSessionID = sessionID + se.currentSession = session + se.currentWorkspace = ws + se.currentSandbox = sb + + return nil +} + +// EndSession ends the current session, cleaning up resources +func (se *SessionExecutor) EndSession(ctx context.Context) error { + if se.currentSessionID == "" { + return nil + } + + // Destroy sandbox if it exists + if se.currentSandbox != nil { + if err := se.currentSandbox.Destroy(ctx); err != nil { + // Log but don't fail + fmt.Printf("Warning: failed to destroy sandbox: %v\n", err) + } + se.currentSandbox = nil + } + + // Release lock + runID := fmt.Sprintf("repl-%d", time.Now().UnixNano()) + se.sessionManager.ReleaseLock(ctx, se.currentSessionID, runID) + + se.currentSessionID = "" + se.currentSession = nil + se.currentWorkspace = nil + + return nil +} + +// SetTools sets the available tools for the session +func (se *SessionExecutor) SetTools(tools []ai.Tool) { + se.tools = tools +} + +// Execute runs a task within the current session, maintaining history +func (se *SessionExecutor) Execute(ctx context.Context, task string) (*harness.ExecutionResult, error) { + if se.currentSessionID == "" { + return nil, fmt.Errorf("no active session - call StartSession first") + } + + // Load existing history + history, err := se.historyStore.Load(se.currentSessionID) + if err != nil { + return nil, fmt.Errorf("failed to load history: %w", err) + } + + // Convert stored messages to ai.Message format for the executor + initialHistory := history.ToAIMessages() + + // Create executor with session state + opts := []harness.ExecutorOption{ + harness.WithWorkspace(se.currentWorkspace), + harness.WithModelName(se.modelName), + } + + if se.systemPrompt != "" { + opts = append(opts, harness.WithSystemPrompt(se.systemPrompt)) + } + + if se.currentSandbox != nil { + opts = append(opts, harness.WithSandbox(se.currentSandbox)) + } + + // Inject history from previous session interactions + if len(initialHistory) > 0 { + opts = append(opts, harness.WithInitialHistory(initialHistory)) + } + + executor := harness.NewAgenticExecutor( + se.genkitApp, + se.harnessConfig, + se.agentConfig, + opts..., + ) + + // Execute with the task + result, err := executor.Execute(ctx, se.currentSessionID, task, se.tools) + + // Save the execution history for session persistence + // Use the full history from the result if available, otherwise construct from task/response + if err == nil && result != nil { + if len(result.History) > 0 { + // Use full history from executor (includes tool calls, etc.) + history.Messages = FromAIMessages(result.History) + } else if result.Success { + // Fallback: append user task and assistant response + newMessages := []StoredMessage{ + { + Role: "user", + Content: task, + Timestamp: time.Now(), + }, + { + Role: "assistant", + Content: result.Response, + Timestamp: time.Now(), + }, + } + history.Messages = append(history.Messages, newMessages...) + } + history.TotalTokens += result.TotalTokens + + if err := se.historyStore.Save(se.currentSessionID, history); err != nil { + // Log but don't fail the execution + fmt.Printf("Warning: failed to save history: %v\n", err) + } + } + + return result, err +} + +// GetHistory returns the current session's message history +func (se *SessionExecutor) GetHistory() (*SessionHistory, error) { + if se.currentSessionID == "" { + return nil, fmt.Errorf("no active session") + } + return se.historyStore.Load(se.currentSessionID) +} + +// ClearHistory clears the current session's message history +func (se *SessionExecutor) ClearHistory() error { + if se.currentSessionID == "" { + return fmt.Errorf("no active session") + } + return se.historyStore.Clear(se.currentSessionID) +} + +// CurrentSession returns the current session info +func (se *SessionExecutor) CurrentSession() *Session { + return se.currentSession +} + +// RefreshLock refreshes the session lock (call periodically for long sessions) +func (se *SessionExecutor) RefreshLock(ctx context.Context) error { + if se.currentSessionID == "" { + return nil + } + runID := fmt.Sprintf("repl-%d", time.Now().UnixNano()) + return se.sessionManager.RefreshLock(ctx, se.currentSessionID, runID) +} diff --git a/pkg/harness/session/session_executor_test.go b/pkg/harness/session/session_executor_test.go new file mode 100644 index 00000000..bdb2cef4 --- /dev/null +++ b/pkg/harness/session/session_executor_test.go @@ -0,0 +1,218 @@ +package session + +import ( + "context" + "os" + "testing" + "time" +) + +func TestHistoryStore_SaveLoad(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "history-test-*") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + defer os.RemoveAll(tmpDir) + + store := NewHistoryStore(tmpDir) + sessionID := "test-session-1" + + // Create session directory + os.MkdirAll(store.historyPath(sessionID)[:len(store.historyPath(sessionID))-len("/.history.json")], 0755) + + // Test empty load (should return empty history) + history, err := store.Load(sessionID) + if err != nil { + t.Fatalf("Failed to load empty history: %v", err) + } + if len(history.Messages) != 0 { + t.Errorf("Expected empty messages, got %d", len(history.Messages)) + } + + // Add some messages + history.Messages = []StoredMessage{ + {Role: "user", Content: "Hello", Timestamp: time.Now()}, + {Role: "assistant", Content: "Hi there!", Timestamp: time.Now()}, + } + history.TotalTokens = 100 + + // Save + if err := store.Save(sessionID, history); err != nil { + t.Fatalf("Failed to save history: %v", err) + } + + // Load again and verify + loaded, err := store.Load(sessionID) + if err != nil { + t.Fatalf("Failed to load history: %v", err) + } + + if len(loaded.Messages) != 2 { + t.Errorf("Expected 2 messages, got %d", len(loaded.Messages)) + } + + if loaded.Messages[0].Content != "Hello" { + t.Errorf("Expected first message 'Hello', got '%s'", loaded.Messages[0].Content) + } + + if loaded.TotalTokens != 100 { + t.Errorf("Expected 100 tokens, got %d", loaded.TotalTokens) + } +} + +func TestHistoryStore_Append(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "history-append-test-*") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + defer os.RemoveAll(tmpDir) + + store := NewHistoryStore(tmpDir) + sessionID := "test-session-2" + + // Append first message + err = store.Append(sessionID, []StoredMessage{ + {Role: "user", Content: "First message", Timestamp: time.Now()}, + }) + if err != nil { + t.Fatalf("Failed to append first message: %v", err) + } + + // Append second message + err = store.Append(sessionID, []StoredMessage{ + {Role: "assistant", Content: "Response", Timestamp: time.Now()}, + }) + if err != nil { + t.Fatalf("Failed to append second message: %v", err) + } + + // Load and verify + history, err := store.Load(sessionID) + if err != nil { + t.Fatalf("Failed to load history: %v", err) + } + + if len(history.Messages) != 2 { + t.Errorf("Expected 2 messages after append, got %d", len(history.Messages)) + } + + if history.Messages[0].Content != "First message" { + t.Errorf("Expected first message 'First message', got '%s'", history.Messages[0].Content) + } +} + +func TestHistoryStore_Clear(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "history-clear-test-*") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + defer os.RemoveAll(tmpDir) + + store := NewHistoryStore(tmpDir) + sessionID := "test-session-3" + + // Create some history + history := &SessionHistory{ + SessionID: sessionID, + Messages: []StoredMessage{ + {Role: "user", Content: "Message", Timestamp: time.Now()}, + }, + } + store.Save(sessionID, history) + + // Clear it + err = store.Clear(sessionID) + if err != nil { + t.Fatalf("Failed to clear history: %v", err) + } + + // Load again - should be empty + loaded, err := store.Load(sessionID) + if err != nil { + t.Fatalf("Failed to load after clear: %v", err) + } + + if len(loaded.Messages) != 0 { + t.Errorf("Expected 0 messages after clear, got %d", len(loaded.Messages)) + } +} + +func TestSessionManager_Integration(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "session-integration-*") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + defer os.RemoveAll(tmpDir) + + ctx := context.Background() + manager := NewManager(tmpDir) + historyStore := NewHistoryStore(tmpDir) + + // Create session + session, err := manager.Create(ctx, "integration-test", "", "") + if err != nil { + t.Fatalf("Failed to create session: %v", err) + } + + // Acquire lock + runID := "test-run-1" + if err := manager.AcquireLock(ctx, session.ID, runID); err != nil { + t.Fatalf("Failed to acquire lock: %v", err) + } + + // Add messages + err = historyStore.Append(session.ID, []StoredMessage{ + {Role: "user", Content: "First input", Timestamp: time.Now()}, + {Role: "assistant", Content: "First response", Timestamp: time.Now()}, + }) + if err != nil { + t.Fatalf("Failed to append messages: %v", err) + } + + // Simulate "closing" the REPL (release lock) + if err := manager.ReleaseLock(ctx, session.ID, runID); err != nil { + t.Fatalf("Failed to release lock: %v", err) + } + + // Simulate "reopening" the REPL (resume session) + resumedSession, err := manager.Get(ctx, session.ID) + if err != nil { + t.Fatalf("Failed to get session: %v", err) + } + + if resumedSession.TotalRuns != 1 { + t.Errorf("Expected TotalRuns=1, got %d", resumedSession.TotalRuns) + } + + // Load history from resumed session + history, err := historyStore.Load(session.ID) + if err != nil { + t.Fatalf("Failed to load history: %v", err) + } + + if len(history.Messages) != 2 { + t.Errorf("Expected 2 messages after resume, got %d", len(history.Messages)) + } + + // Add more messages in new "session" + runID2 := "test-run-2" + if err := manager.AcquireLock(ctx, session.ID, runID2); err != nil { + t.Fatalf("Failed to acquire lock for run 2: %v", err) + } + + err = historyStore.Append(session.ID, []StoredMessage{ + {Role: "user", Content: "Second input", Timestamp: time.Now()}, + {Role: "assistant", Content: "Second response", Timestamp: time.Now()}, + }) + if err != nil { + t.Fatalf("Failed to append more messages: %v", err) + } + + // Verify all 4 messages are present + history, _ = historyStore.Load(session.ID) + if len(history.Messages) != 4 { + t.Errorf("Expected 4 messages total, got %d", len(history.Messages)) + } + + t.Logf("Session integration test passed - history persists across REPL sessions") +} diff --git a/pkg/harness/skills/fs_backend.go b/pkg/harness/skills/fs_backend.go new file mode 100644 index 00000000..1a4cdab5 --- /dev/null +++ b/pkg/harness/skills/fs_backend.go @@ -0,0 +1,36 @@ +package skills + +import ( + "os" + "path/filepath" +) + +// FSBackend implements Backend using the local filesystem +type FSBackend struct{} + +// NewFSBackend creates a new filesystem backend +func NewFSBackend() *FSBackend { + return &FSBackend{} +} + +// ListDir lists directory entries +func (b *FSBackend) ListDir(path string) ([]DirEntry, error) { + entries, err := os.ReadDir(path) + if err != nil { + return nil, err + } + + result := make([]DirEntry, len(entries)) + for i, e := range entries { + result[i] = DirEntry{ + Name: e.Name(), + IsDir: e.IsDir(), + } + } + return result, nil +} + +// ReadFile reads a file's content +func (b *FSBackend) ReadFile(path string) ([]byte, error) { + return os.ReadFile(filepath.Clean(path)) +} diff --git a/pkg/harness/skills/middleware.go b/pkg/harness/skills/middleware.go new file mode 100644 index 00000000..b7f1a80b --- /dev/null +++ b/pkg/harness/skills/middleware.go @@ -0,0 +1,164 @@ +package skills + +import ( + "fmt" + "log/slog" + "os" + "path/filepath" + "strings" +) + +// SkillsMiddleware loads and injects skills into the system prompt +type SkillsMiddleware struct { + backend Backend + sources []string + logger *slog.Logger +} + +// NewSkillsMiddleware creates a new skills middleware with the given backend and sources +func NewSkillsMiddleware(backend Backend, sources []string) *SkillsMiddleware { + return &SkillsMiddleware{ + backend: backend, + sources: sources, + logger: slog.Default().With("component", "skills_middleware"), + } +} + +// LoadSkills discovers all skills from configured sources +// Later sources override earlier ones (project > user) +func (m *SkillsMiddleware) LoadSkills() ([]SkillMetadata, error) { + allSkills := make(map[string]SkillMetadata) + + for _, source := range m.sources { + expandedPath := expandPath(source) + skills, err := m.loadSkillsFromSource(expandedPath) + if err != nil { + // Optional sources may not exist - just log and continue + m.logger.Debug("skill source not found or error", + "source", expandedPath, + "error", err) + continue + } + for _, skill := range skills { + allSkills[skill.Name] = skill // Later wins + } + } + + result := make([]SkillMetadata, 0, len(allSkills)) + for _, skill := range allSkills { + result = append(result, skill) + } + return result, nil +} + +func (m *SkillsMiddleware) loadSkillsFromSource(basePath string) ([]SkillMetadata, error) { + var skills []SkillMetadata + + entries, err := m.backend.ListDir(basePath) + if err != nil { + return nil, fmt.Errorf("failed to list skills directory %s: %w", basePath, err) + } + + for _, entry := range entries { + if !entry.IsDir { + continue + } + + skillPath := filepath.Join(basePath, entry.Name, "SKILL.md") + content, err := m.backend.ReadFile(skillPath) + if err != nil { + // Skill directories without SKILL.md are skipped + m.logger.Debug("SKILL.md not found", + "path", skillPath, + "error", err) + continue + } + + metadata, err := ParseSkillMetadata(content, skillPath, entry.Name) + if err != nil { + m.logger.Warn("failed to parse skill", + "path", skillPath, + "error", err) + continue + } + + skills = append(skills, *metadata) + m.logger.Debug("loaded skill", + "name", metadata.Name, + "path", skillPath) + } + + return skills, nil +} + +// FormatSystemPromptSection generates the skills section for the system prompt +// This implements progressive disclosure - only skill names and descriptions +// are shown; the agent must read SKILL.md files for full instructions +func (m *SkillsMiddleware) FormatSystemPromptSection(skills []SkillMetadata) string { + if len(skills) == 0 { + return "" + } + + var sb strings.Builder + sb.WriteString("\n\n## Skills System\n\n") + sb.WriteString("You have access to specialized skills that provide domain-specific workflows.\n\n") + + // Format available skills + sb.WriteString("**Available Skills:**\n\n") + for _, skill := range skills { + sb.WriteString(fmt.Sprintf("- **%s**: %s\n", skill.Name, skill.Description)) + if len(skill.AllowedTools) > 0 { + sb.WriteString(fmt.Sprintf(" → Allowed tools: %s\n", + strings.Join(skill.AllowedTools, ", "))) + } + sb.WriteString(fmt.Sprintf(" → Read `%s` for full instructions\n", skill.Path)) + } + + // Add usage instructions + sb.WriteString(` +**How to Use Skills (Progressive Disclosure):** + +1. **Recognize when a skill applies**: Check if the task matches a skill's description +2. **Read the skill's full instructions**: Use read_file with the path above +3. **Follow the skill's workflow**: SKILL.md contains step-by-step instructions +4. **Only read when needed**: Don't read skills you won't use (saves tokens) + +**When to Use Skills:** +- User's request matches a skill's domain +- You need specialized knowledge or structured workflows +- A skill provides proven patterns for complex tasks +`) + + return sb.String() +} + +// GetSkillInstructions reads the full instructions from a skill file +func (m *SkillsMiddleware) GetSkillInstructions(skillPath string) (string, error) { + content, err := m.backend.ReadFile(skillPath) + if err != nil { + return "", fmt.Errorf("failed to read skill file: %w", err) + } + + return GetSkillBody(content), nil +} + +// expandPath expands ~ to the user's home directory +func expandPath(path string) string { + if strings.HasPrefix(path, "~/") { + home, err := os.UserHomeDir() + if err != nil { + return path + } + return filepath.Join(home, path[2:]) + } + return path +} + +// DefaultSkillSources returns the default skill source paths +func DefaultSkillSources(envPath string) []string { + return []string{ + filepath.Join(envPath, "skills", "user"), + filepath.Join(envPath, "skills", "project"), + ".station/skills", + } +} diff --git a/pkg/harness/skills/middleware_test.go b/pkg/harness/skills/middleware_test.go new file mode 100644 index 00000000..8e497907 --- /dev/null +++ b/pkg/harness/skills/middleware_test.go @@ -0,0 +1,313 @@ +package skills + +import ( + "strings" + "testing" +) + +// MockBackend implements Backend for testing +type MockBackend struct { + Files map[string][]byte + Dirs map[string][]DirEntry +} + +func NewMockBackend() *MockBackend { + return &MockBackend{ + Files: make(map[string][]byte), + Dirs: make(map[string][]DirEntry), + } +} + +func (m *MockBackend) ListDir(path string) ([]DirEntry, error) { + if entries, ok := m.Dirs[path]; ok { + return entries, nil + } + return nil, &mockError{msg: "directory not found: " + path} +} + +func (m *MockBackend) ReadFile(path string) ([]byte, error) { + if content, ok := m.Files[path]; ok { + return content, nil + } + return nil, &mockError{msg: "file not found: " + path} +} + +type mockError struct { + msg string +} + +func (e *mockError) Error() string { + return e.msg +} + +func TestParseSkillMetadata(t *testing.T) { + tests := []struct { + name string + content string + dirName string + wantName string + wantDesc string + wantErr bool + }{ + { + name: "valid skill with all fields", + content: `--- +name: web-research +description: Structured approach to web research +license: MIT +allowed-tools: + - web_search + - web_fetch +triggers: + - research + - search +--- + +# Web Research Skill + +Instructions here...`, + dirName: "web-research", + wantName: "web-research", + wantDesc: "Structured approach to web research", + wantErr: false, + }, + { + name: "skill without name uses directory", + content: `--- +description: Test skill +--- + +# Test Skill`, + dirName: "test-skill", + wantName: "test-skill", + wantDesc: "Test skill", + wantErr: false, + }, + { + name: "name mismatch error", + content: `--- +name: wrong-name +description: Test +--- + +Content`, + dirName: "actual-name", + wantErr: true, + }, + { + name: "no frontmatter", + content: "# Just markdown content", + dirName: "test", + wantErr: true, + }, + { + name: "unclosed frontmatter", + content: "---\nname: test", + dirName: "test", + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + metadata, err := ParseSkillMetadata([]byte(tt.content), "/test/path/SKILL.md", tt.dirName) + if (err != nil) != tt.wantErr { + t.Errorf("ParseSkillMetadata() error = %v, wantErr %v", err, tt.wantErr) + return + } + if err != nil { + return + } + if metadata.Name != tt.wantName { + t.Errorf("Name = %q, want %q", metadata.Name, tt.wantName) + } + if metadata.Description != tt.wantDesc { + t.Errorf("Description = %q, want %q", metadata.Description, tt.wantDesc) + } + }) + } +} + +func TestGetSkillBody(t *testing.T) { + content := `--- +name: test +description: Test skill +--- + +# Test Skill + +This is the body content.` + + body := GetSkillBody([]byte(content)) + if !strings.Contains(body, "# Test Skill") { + t.Error("body should contain heading") + } + if !strings.Contains(body, "This is the body content.") { + t.Error("body should contain content") + } + if strings.Contains(body, "name: test") { + t.Error("body should not contain frontmatter") + } +} + +func TestSkillsMiddleware_LoadSkills(t *testing.T) { + backend := NewMockBackend() + + // Set up mock filesystem + backend.Dirs["/skills/user"] = []DirEntry{ + {Name: "web-research", IsDir: true}, + {Name: "code-review", IsDir: true}, + {Name: "not-a-skill.txt", IsDir: false}, // Should be ignored + } + + backend.Files["/skills/user/web-research/SKILL.md"] = []byte(`--- +name: web-research +description: Web research skill +--- +# Web Research`) + + backend.Files["/skills/user/code-review/SKILL.md"] = []byte(`--- +name: code-review +description: Code review skill +--- +# Code Review`) + + backend.Dirs["/skills/project"] = []DirEntry{ + {Name: "code-review", IsDir: true}, // Override user skill + } + + backend.Files["/skills/project/code-review/SKILL.md"] = []byte(`--- +name: code-review +description: Project-specific code review +--- +# Project Code Review`) + + middleware := NewSkillsMiddleware(backend, []string{"/skills/user", "/skills/project"}) + skills, err := middleware.LoadSkills() + if err != nil { + t.Fatalf("LoadSkills() error = %v", err) + } + + if len(skills) != 2 { + t.Errorf("LoadSkills() returned %d skills, want 2", len(skills)) + } + + // Check that code-review was overridden by project + for _, skill := range skills { + if skill.Name == "code-review" { + if skill.Description != "Project-specific code review" { + t.Errorf("code-review should have project description, got %q", skill.Description) + } + } + } +} + +func TestSkillsMiddleware_FormatSystemPromptSection(t *testing.T) { + skills := []SkillMetadata{ + { + Name: "web-research", + Description: "Web research skill", + AllowedTools: []string{"web_search", "web_fetch"}, + Path: "/skills/user/web-research/SKILL.md", + }, + { + Name: "code-review", + Description: "Code review skill", + Path: "/skills/user/code-review/SKILL.md", + }, + } + + middleware := NewSkillsMiddleware(nil, nil) + section := middleware.FormatSystemPromptSection(skills) + + // Check key elements + if !strings.Contains(section, "## Skills System") { + t.Error("section should contain header") + } + if !strings.Contains(section, "**web-research**") { + t.Error("section should contain skill name") + } + if !strings.Contains(section, "Web research skill") { + t.Error("section should contain skill description") + } + if !strings.Contains(section, "web_search, web_fetch") { + t.Error("section should contain allowed tools") + } + if !strings.Contains(section, "/skills/user/web-research/SKILL.md") { + t.Error("section should contain skill path") + } + if !strings.Contains(section, "Progressive Disclosure") { + t.Error("section should contain usage instructions") + } +} + +func TestSkillsMiddleware_EmptySkills(t *testing.T) { + middleware := NewSkillsMiddleware(nil, nil) + section := middleware.FormatSystemPromptSection(nil) + if section != "" { + t.Errorf("empty skills should produce empty section, got %q", section) + } +} + +func TestExpandPath(t *testing.T) { + tests := []struct { + input string + want string + }{ + {"/absolute/path", "/absolute/path"}, + {"relative/path", "relative/path"}, + // Note: can't easily test ~ expansion without mocking os.UserHomeDir + } + + for _, tt := range tests { + t.Run(tt.input, func(t *testing.T) { + got := expandPath(tt.input) + if !strings.HasPrefix(tt.input, "~/") && got != tt.want { + t.Errorf("expandPath(%q) = %q, want %q", tt.input, got, tt.want) + } + }) + } +} + +func TestSkillsMiddleware_MissingSource(t *testing.T) { + backend := NewMockBackend() + // No directories set up - all sources should fail gracefully + + middleware := NewSkillsMiddleware(backend, []string{"/nonexistent/path"}) + skills, err := middleware.LoadSkills() + if err != nil { + t.Fatalf("LoadSkills() should not error on missing sources, got %v", err) + } + if len(skills) != 0 { + t.Errorf("LoadSkills() should return empty for missing sources, got %d skills", len(skills)) + } +} + +func TestSkillsMiddleware_SkillWithoutSkillMd(t *testing.T) { + backend := NewMockBackend() + backend.Dirs["/skills"] = []DirEntry{ + {Name: "valid-skill", IsDir: true}, + {Name: "no-skill-md", IsDir: true}, + } + + backend.Files["/skills/valid-skill/SKILL.md"] = []byte(`--- +name: valid-skill +description: Valid skill +--- +# Valid`) + + // no-skill-md directory has no SKILL.md file + + middleware := NewSkillsMiddleware(backend, []string{"/skills"}) + skills, err := middleware.LoadSkills() + if err != nil { + t.Fatalf("LoadSkills() error = %v", err) + } + + if len(skills) != 1 { + t.Errorf("LoadSkills() returned %d skills, want 1", len(skills)) + } + if skills[0].Name != "valid-skill" { + t.Errorf("skill name = %q, want valid-skill", skills[0].Name) + } +} diff --git a/pkg/harness/skills/parser.go b/pkg/harness/skills/parser.go new file mode 100644 index 00000000..f1b0eceb --- /dev/null +++ b/pkg/harness/skills/parser.go @@ -0,0 +1,62 @@ +package skills + +import ( + "fmt" + "strings" + + "gopkg.in/yaml.v3" +) + +// ParseSkillMetadata extracts YAML frontmatter from SKILL.md content +func ParseSkillMetadata(content []byte, path, dirName string) (*SkillMetadata, error) { + contentStr := string(content) + + // Extract YAML frontmatter between --- markers + if !strings.HasPrefix(contentStr, "---") { + return nil, fmt.Errorf("no frontmatter found in %s", path) + } + + // Find the closing --- + rest := contentStr[3:] + parts := strings.SplitN(rest, "---", 2) + if len(parts) < 2 { + return nil, fmt.Errorf("frontmatter not closed in %s", path) + } + + frontmatter := strings.TrimSpace(parts[0]) + + var metadata SkillMetadata + if err := yaml.Unmarshal([]byte(frontmatter), &metadata); err != nil { + return nil, fmt.Errorf("failed to parse frontmatter in %s: %w", path, err) + } + + // Validate name matches directory if name is specified + if metadata.Name != "" && metadata.Name != dirName { + return nil, fmt.Errorf("skill name %q must match directory %q in %s", metadata.Name, dirName, path) + } + + // Use directory name if name not specified + if metadata.Name == "" { + metadata.Name = dirName + } + + metadata.Path = path + return &metadata, nil +} + +// GetSkillBody returns the content after the frontmatter (the actual instructions) +func GetSkillBody(content []byte) string { + contentStr := string(content) + + if !strings.HasPrefix(contentStr, "---") { + return contentStr + } + + rest := contentStr[3:] + parts := strings.SplitN(rest, "---", 2) + if len(parts) < 2 { + return contentStr + } + + return strings.TrimSpace(parts[1]) +} diff --git a/pkg/harness/skills/types.go b/pkg/harness/skills/types.go new file mode 100644 index 00000000..299b63fe --- /dev/null +++ b/pkg/harness/skills/types.go @@ -0,0 +1,28 @@ +package skills + +// SkillMetadata represents parsed SKILL.md frontmatter +type SkillMetadata struct { + Name string `yaml:"name"` + Description string `yaml:"description"` + License string `yaml:"license,omitempty"` + AllowedTools []string `yaml:"allowed-tools,omitempty"` + Triggers []string `yaml:"triggers,omitempty"` + Path string `yaml:"-"` // Set after parsing, not from YAML +} + +// SkillsConfig configures the skills middleware from agent frontmatter +type SkillsConfig struct { + Sources []string `yaml:"sources,omitempty" json:"sources,omitempty"` +} + +// DirEntry represents a directory entry from the backend +type DirEntry struct { + Name string + IsDir bool +} + +// Backend interface for file operations (can be workspace, sandbox, or direct FS) +type Backend interface { + ListDir(path string) ([]DirEntry, error) + ReadFile(path string) ([]byte, error) +} diff --git a/pkg/harness/stream/nats_publisher.go b/pkg/harness/stream/nats_publisher.go new file mode 100644 index 00000000..f3dda786 --- /dev/null +++ b/pkg/harness/stream/nats_publisher.go @@ -0,0 +1,126 @@ +package stream + +import ( + "context" + "encoding/json" + "fmt" + + "github.com/nats-io/nats.go" +) + +type NATSPublisher struct { + nc *nats.Conn + js nats.JetStreamContext + stationID string + useJS bool +} + +type NATSPublisherConfig struct { + StationID string + UseJetStream bool +} + +func NewNATSPublisher(nc *nats.Conn, cfg NATSPublisherConfig) (*NATSPublisher, error) { + p := &NATSPublisher{ + nc: nc, + stationID: cfg.StationID, + useJS: cfg.UseJetStream, + } + + if cfg.UseJetStream { + js, err := nc.JetStream() + if err != nil { + return nil, fmt.Errorf("failed to get JetStream context: %w", err) + } + p.js = js + } + + return p, nil +} + +func (p *NATSPublisher) subject(runID string) string { + return fmt.Sprintf("station.%s.run.%s.stream", p.stationID, runID) +} + +func (p *NATSPublisher) Publish(ctx context.Context, event *Event) error { + data, err := event.JSON() + if err != nil { + return fmt.Errorf("failed to marshal event: %w", err) + } + + subject := p.subject(event.RunUUID) + + if p.useJS && p.js != nil { + _, err = p.js.Publish(subject, data) + } else { + err = p.nc.Publish(subject, data) + } + + if err != nil { + return fmt.Errorf("failed to publish to %s: %w", subject, err) + } + + return nil +} + +func (p *NATSPublisher) Close() error { + return nil +} + +type NATSSubscriber struct { + nc *nats.Conn + stationID string +} + +func NewNATSSubscriber(nc *nats.Conn, stationID string) *NATSSubscriber { + return &NATSSubscriber{ + nc: nc, + stationID: stationID, + } +} + +func (s *NATSSubscriber) SubscribeRun(ctx context.Context, runID string, handler func(*Event)) (*nats.Subscription, error) { + subject := fmt.Sprintf("station.%s.run.%s.stream", s.stationID, runID) + + sub, err := s.nc.Subscribe(subject, func(msg *nats.Msg) { + var event Event + if err := json.Unmarshal(msg.Data, &event); err != nil { + return + } + handler(&event) + }) + + if err != nil { + return nil, fmt.Errorf("failed to subscribe to %s: %w", subject, err) + } + + go func() { + <-ctx.Done() + sub.Unsubscribe() + }() + + return sub, nil +} + +func (s *NATSSubscriber) SubscribeStation(ctx context.Context, handler func(*Event)) (*nats.Subscription, error) { + subject := fmt.Sprintf("station.%s.run.*.stream", s.stationID) + + sub, err := s.nc.Subscribe(subject, func(msg *nats.Msg) { + var event Event + if err := json.Unmarshal(msg.Data, &event); err != nil { + return + } + handler(&event) + }) + + if err != nil { + return nil, fmt.Errorf("failed to subscribe to %s: %w", subject, err) + } + + go func() { + <-ctx.Done() + sub.Unsubscribe() + }() + + return sub, nil +} diff --git a/pkg/harness/stream/publisher.go b/pkg/harness/stream/publisher.go new file mode 100644 index 00000000..1e9183b7 --- /dev/null +++ b/pkg/harness/stream/publisher.go @@ -0,0 +1,277 @@ +package stream + +import ( + "context" + "encoding/json" + "sync/atomic" + "time" +) + +type EventType string + +const ( + EventToken EventType = "token" + EventThinking EventType = "thinking" + EventToolStart EventType = "tool_start" + EventToolResult EventType = "tool_result" + EventStepComplete EventType = "step_complete" + EventRunStart EventType = "run_start" + EventRunComplete EventType = "run_complete" + EventError EventType = "error" +) + +type Event struct { + StationRunID string `json:"station_run_id"` + RunUUID string `json:"run_uuid"` + WorkflowRunID string `json:"workflow_run_id,omitempty"` + SessionID string `json:"session_id,omitempty"` + AgentID string `json:"agent_id"` + AgentName string `json:"agent_name,omitempty"` + StationID string `json:"station_id,omitempty"` + Seq int64 `json:"seq"` + Timestamp time.Time `json:"timestamp"` + Type EventType `json:"type"` + Data any `json:"data"` +} + +type TokenData struct { + Content string `json:"content"` + Done bool `json:"done"` +} + +type ThinkingData struct { + Content string `json:"content"` +} + +type ToolStartData struct { + ToolName string `json:"tool_name"` + ToolID string `json:"tool_id"` + Input any `json:"input"` +} + +type ToolResultData struct { + ToolName string `json:"tool_name"` + ToolID string `json:"tool_id"` + Output any `json:"output"` + DurationMs int64 `json:"duration_ms"` + Error string `json:"error,omitempty"` +} + +type StepCompleteData struct { + StepNumber int `json:"step_number"` + TotalTokens int `json:"total_tokens"` + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` + FinishReason string `json:"finish_reason"` +} + +type RunStartData struct { + AgentID string `json:"agent_id"` + AgentName string `json:"agent_name"` + Task string `json:"task"` + MaxSteps int `json:"max_steps"` +} + +type RunCompleteData struct { + Success bool `json:"success"` + TotalSteps int `json:"total_steps"` + TotalTokens int `json:"total_tokens"` + DurationMs int64 `json:"duration_ms"` + FinishReason string `json:"finish_reason"` + Error string `json:"error,omitempty"` +} + +type Publisher interface { + Publish(ctx context.Context, event *Event) error + Close() error +} + +type StreamIdentifiers struct { + StationRunID string + RunUUID string + WorkflowRunID string + SessionID string + AgentID string + AgentName string + StationID string +} + +type StreamContext struct { + ids StreamIdentifiers + seq int64 + publisher Publisher +} + +func NewStreamContext(ids StreamIdentifiers, publisher Publisher) *StreamContext { + return &StreamContext{ + ids: ids, + seq: 0, + publisher: publisher, + } +} + +func (s *StreamContext) nextSeq() int64 { + return atomic.AddInt64(&s.seq, 1) +} + +func (s *StreamContext) emit(ctx context.Context, eventType EventType, data any) error { + if s.publisher == nil { + return nil + } + + event := &Event{ + StationRunID: s.ids.StationRunID, + RunUUID: s.ids.RunUUID, + WorkflowRunID: s.ids.WorkflowRunID, + SessionID: s.ids.SessionID, + AgentID: s.ids.AgentID, + AgentName: s.ids.AgentName, + StationID: s.ids.StationID, + Seq: s.nextSeq(), + Timestamp: time.Now(), + Type: eventType, + Data: data, + } + + return s.publisher.Publish(ctx, event) +} + +func (s *StreamContext) Identifiers() StreamIdentifiers { + return s.ids +} + +func (s *StreamContext) EmitRunStart(ctx context.Context, agentID, agentName, task string, maxSteps int) error { + return s.emit(ctx, EventRunStart, RunStartData{ + AgentID: agentID, + AgentName: agentName, + Task: task, + MaxSteps: maxSteps, + }) +} + +func (s *StreamContext) EmitToken(ctx context.Context, content string, done bool) error { + return s.emit(ctx, EventToken, TokenData{ + Content: content, + Done: done, + }) +} + +func (s *StreamContext) EmitThinking(ctx context.Context, content string) error { + return s.emit(ctx, EventThinking, ThinkingData{ + Content: content, + }) +} + +func (s *StreamContext) EmitToolStart(ctx context.Context, toolName, toolID string, input any) error { + return s.emit(ctx, EventToolStart, ToolStartData{ + ToolName: toolName, + ToolID: toolID, + Input: input, + }) +} + +func (s *StreamContext) EmitToolResult(ctx context.Context, toolName, toolID string, output any, durationMs int64, err error) error { + errStr := "" + if err != nil { + errStr = err.Error() + } + return s.emit(ctx, EventToolResult, ToolResultData{ + ToolName: toolName, + ToolID: toolID, + Output: output, + DurationMs: durationMs, + Error: errStr, + }) +} + +func (s *StreamContext) EmitStepComplete(ctx context.Context, stepNum, totalTokens, inputTokens, outputTokens int, finishReason string) error { + return s.emit(ctx, EventStepComplete, StepCompleteData{ + StepNumber: stepNum, + TotalTokens: totalTokens, + InputTokens: inputTokens, + OutputTokens: outputTokens, + FinishReason: finishReason, + }) +} + +func (s *StreamContext) EmitRunComplete(ctx context.Context, success bool, steps, tokens int, durationMs int64, finishReason string, err error) error { + errStr := "" + if err != nil { + errStr = err.Error() + } + return s.emit(ctx, EventRunComplete, RunCompleteData{ + Success: success, + TotalSteps: steps, + TotalTokens: tokens, + DurationMs: durationMs, + FinishReason: finishReason, + Error: errStr, + }) +} + +func (s *StreamContext) EmitError(ctx context.Context, err error) error { + return s.emit(ctx, EventError, map[string]string{"error": err.Error()}) +} + +type MultiPublisher struct { + publishers []Publisher +} + +func NewMultiPublisher(publishers ...Publisher) *MultiPublisher { + return &MultiPublisher{publishers: publishers} +} + +func (m *MultiPublisher) Publish(ctx context.Context, event *Event) error { + for _, p := range m.publishers { + if err := p.Publish(ctx, event); err != nil { + return err + } + } + return nil +} + +func (m *MultiPublisher) Close() error { + for _, p := range m.publishers { + p.Close() + } + return nil +} + +type NoOpPublisher struct{} + +func (n *NoOpPublisher) Publish(ctx context.Context, event *Event) error { return nil } +func (n *NoOpPublisher) Close() error { return nil } + +type ChannelPublisher struct { + ch chan *Event +} + +func NewChannelPublisher(bufferSize int) *ChannelPublisher { + return &ChannelPublisher{ + ch: make(chan *Event, bufferSize), + } +} + +func (c *ChannelPublisher) Publish(ctx context.Context, event *Event) error { + select { + case c.ch <- event: + return nil + case <-ctx.Done(): + return ctx.Err() + default: + return nil + } +} + +func (c *ChannelPublisher) Events() <-chan *Event { + return c.ch +} + +func (c *ChannelPublisher) Close() error { + close(c.ch) + return nil +} + +func (e *Event) JSON() ([]byte, error) { + return json.Marshal(e) +} diff --git a/pkg/harness/stream/publisher_test.go b/pkg/harness/stream/publisher_test.go new file mode 100644 index 00000000..88a38310 --- /dev/null +++ b/pkg/harness/stream/publisher_test.go @@ -0,0 +1,261 @@ +package stream + +import ( + "context" + "testing" + "time" +) + +func testIdentifiers() StreamIdentifiers { + return StreamIdentifiers{ + StationRunID: "123", + RunUUID: "run-uuid-123", + WorkflowRunID: "wf-456", + SessionID: "session-456", + AgentID: "agent-1", + AgentName: "test-agent", + StationID: "station-abc", + } +} + +func TestEvent_JSON(t *testing.T) { + event := &Event{ + StationRunID: "123", + RunUUID: "run-uuid-123", + WorkflowRunID: "wf-456", + SessionID: "session-456", + AgentID: "agent-1", + AgentName: "test-agent", + StationID: "station-abc", + Seq: 1, + Timestamp: time.Now(), + Type: EventToken, + Data: TokenData{ + Content: "Hello", + Done: false, + }, + } + + data, err := event.JSON() + if err != nil { + t.Fatalf("expected no error, got %v", err) + } + + if len(data) == 0 { + t.Fatal("expected non-empty JSON") + } +} + +func TestStreamContext_EmitToken(t *testing.T) { + pub := NewChannelPublisher(10) + ctx := context.Background() + ids := testIdentifiers() + sc := NewStreamContext(ids, pub) + + err := sc.EmitToken(ctx, "Hello", false) + if err != nil { + t.Fatalf("expected no error, got %v", err) + } + + select { + case event := <-pub.Events(): + if event.Type != EventToken { + t.Errorf("expected EventToken, got %v", event.Type) + } + if event.RunUUID != "run-uuid-123" { + t.Errorf("expected run-uuid-123, got %s", event.RunUUID) + } + if event.StationRunID != "123" { + t.Errorf("expected 123, got %s", event.StationRunID) + } + if event.WorkflowRunID != "wf-456" { + t.Errorf("expected wf-456, got %s", event.WorkflowRunID) + } + if event.StationID != "station-abc" { + t.Errorf("expected station-abc, got %s", event.StationID) + } + if event.AgentID != "agent-1" { + t.Errorf("expected agent-1, got %s", event.AgentID) + } + if event.Seq != 1 { + t.Errorf("expected seq 1, got %d", event.Seq) + } + case <-time.After(time.Second): + t.Fatal("timeout waiting for event") + } +} + +func TestStreamContext_EmitToolStart(t *testing.T) { + pub := NewChannelPublisher(10) + ctx := context.Background() + sc := NewStreamContext(testIdentifiers(), pub) + + err := sc.EmitToolStart(ctx, "bash", "tool-1", map[string]string{"command": "ls"}) + if err != nil { + t.Fatalf("expected no error, got %v", err) + } + + select { + case event := <-pub.Events(): + if event.Type != EventToolStart { + t.Errorf("expected EventToolStart, got %v", event.Type) + } + data, ok := event.Data.(ToolStartData) + if !ok { + t.Fatal("expected ToolStartData") + } + if data.ToolName != "bash" { + t.Errorf("expected bash, got %s", data.ToolName) + } + case <-time.After(time.Second): + t.Fatal("timeout waiting for event") + } +} + +func TestStreamContext_EmitRunComplete(t *testing.T) { + pub := NewChannelPublisher(10) + ctx := context.Background() + sc := NewStreamContext(testIdentifiers(), pub) + + err := sc.EmitRunComplete(ctx, true, 5, 1000, 5000, "agent_done", nil) + if err != nil { + t.Fatalf("expected no error, got %v", err) + } + + select { + case event := <-pub.Events(): + if event.Type != EventRunComplete { + t.Errorf("expected EventRunComplete, got %v", event.Type) + } + data, ok := event.Data.(RunCompleteData) + if !ok { + t.Fatal("expected RunCompleteData") + } + if !data.Success { + t.Error("expected Success=true") + } + if data.TotalSteps != 5 { + t.Errorf("expected 5 steps, got %d", data.TotalSteps) + } + case <-time.After(time.Second): + t.Fatal("timeout waiting for event") + } +} + +func TestNoOpPublisher(t *testing.T) { + pub := &NoOpPublisher{} + ctx := context.Background() + + err := pub.Publish(ctx, &Event{Type: EventToken}) + if err != nil { + t.Fatalf("expected no error, got %v", err) + } + + err = pub.Close() + if err != nil { + t.Fatalf("expected no error, got %v", err) + } +} + +func TestMultiPublisher(t *testing.T) { + pub1 := NewChannelPublisher(10) + pub2 := NewChannelPublisher(10) + multi := NewMultiPublisher(pub1, pub2) + ctx := context.Background() + + event := &Event{ + RunUUID: "run-uuid-123", + StationRunID: "123", + Type: EventToken, + Data: TokenData{Content: "test"}, + } + + err := multi.Publish(ctx, event) + if err != nil { + t.Fatalf("expected no error, got %v", err) + } + + select { + case e := <-pub1.Events(): + if e.RunUUID != "run-uuid-123" { + t.Errorf("pub1: expected run-uuid-123, got %s", e.RunUUID) + } + case <-time.After(time.Second): + t.Fatal("timeout waiting for pub1 event") + } + + select { + case e := <-pub2.Events(): + if e.RunUUID != "run-uuid-123" { + t.Errorf("pub2: expected run-uuid-123, got %s", e.RunUUID) + } + case <-time.After(time.Second): + t.Fatal("timeout waiting for pub2 event") + } +} + +func TestChannelPublisher_ContextCancellation(t *testing.T) { + pub := NewChannelPublisher(1) + ctx, cancel := context.WithCancel(context.Background()) + + pub.Publish(ctx, &Event{Type: EventToken}) + + cancel() + err := pub.Publish(ctx, &Event{Type: EventToken}) + if err != context.Canceled { + t.Errorf("expected context.Canceled error, got %v", err) + } +} + +func TestStreamContext_NilPublisher(t *testing.T) { + sc := NewStreamContext(testIdentifiers(), nil) + ctx := context.Background() + + err := sc.EmitToken(ctx, "test", false) + if err != nil { + t.Fatalf("expected no error with nil publisher, got %v", err) + } +} + +func TestStreamContext_SequenceIncrement(t *testing.T) { + pub := NewChannelPublisher(10) + ctx := context.Background() + sc := NewStreamContext(testIdentifiers(), pub) + + sc.EmitToken(ctx, "a", false) + sc.EmitToken(ctx, "b", false) + sc.EmitToken(ctx, "c", true) + + var seqs []int64 + for i := 0; i < 3; i++ { + select { + case event := <-pub.Events(): + seqs = append(seqs, event.Seq) + case <-time.After(time.Second): + t.Fatal("timeout") + } + } + + if seqs[0] != 1 || seqs[1] != 2 || seqs[2] != 3 { + t.Errorf("expected sequences [1,2,3], got %v", seqs) + } +} + +func TestStreamContext_Identifiers(t *testing.T) { + ids := testIdentifiers() + sc := NewStreamContext(ids, nil) + + got := sc.Identifiers() + if got.StationRunID != ids.StationRunID { + t.Errorf("StationRunID mismatch: %s vs %s", got.StationRunID, ids.StationRunID) + } + if got.RunUUID != ids.RunUUID { + t.Errorf("RunUUID mismatch: %s vs %s", got.RunUUID, ids.RunUUID) + } + if got.WorkflowRunID != ids.WorkflowRunID { + t.Errorf("WorkflowRunID mismatch: %s vs %s", got.WorkflowRunID, ids.WorkflowRunID) + } + if got.StationID != ids.StationID { + t.Errorf("StationID mismatch: %s vs %s", got.StationID, ids.StationID) + } +} diff --git a/pkg/harness/tools/bash.go b/pkg/harness/tools/bash.go new file mode 100644 index 00000000..1787ed06 --- /dev/null +++ b/pkg/harness/tools/bash.go @@ -0,0 +1,188 @@ +package tools + +import ( + "bytes" + "context" + "fmt" + "os/exec" + "path/filepath" + "strings" + "time" + + "station/pkg/harness/sandbox" + + "github.com/firebase/genkit/go/ai" + "github.com/firebase/genkit/go/genkit" +) + +const ( + defaultBashTimeout = 2 * time.Minute + maxBashOutputLength = 30000 +) + +type BashInput struct { + Command string `json:"command" jsonschema:"description=The command to execute"` + Timeout int `json:"timeout,omitempty" jsonschema:"description=Optional timeout in milliseconds (default 120000)"` + Workdir string `json:"workdir,omitempty" jsonschema:"description=Working directory for the command. Defaults to workspace root."` + Description string `json:"description" jsonschema:"description=Clear 5-10 word description of what this command does"` +} + +type BashOutput struct { + Output string `json:"output"` + ExitCode int `json:"exit_code"` + TimedOut bool `json:"timed_out,omitempty"` +} + +// NewBashTool creates a bash tool that uses the host directly (legacy API). +// Prefer NewBashToolWithSandbox for sandbox isolation support. +func NewBashTool(genkitApp *genkit.Genkit, workspacePath string) ai.Tool { + return NewBashToolWithSandbox(genkitApp, workspacePath, nil) +} + +// NewBashToolWithSandbox creates a bash tool with optional sandbox support. +// If sb is nil, commands are executed directly on the host. +func NewBashToolWithSandbox(genkitApp *genkit.Genkit, workspacePath string, sb sandbox.Sandbox) ai.Tool { + return genkit.DefineTool( + genkitApp, + "bash", + `Execute a bash command in a persistent shell session. + +IMPORTANT: +- Use the workdir parameter instead of 'cd' commands +- Quote file paths containing spaces with double quotes +- Commands timeout after 2 minutes by default +- Output is truncated at 30000 characters + +Examples: +- ls -la +- git status +- npm install +- mkdir "path with spaces"`, + func(ctx *ai.ToolContext, input BashInput) (BashOutput, error) { + if input.Command == "" { + return BashOutput{}, fmt.Errorf("command is required") + } + + workdir := workspacePath + if input.Workdir != "" { + if filepath.IsAbs(input.Workdir) { + workdir = input.Workdir + } else { + workdir = filepath.Join(workspacePath, input.Workdir) + } + } + + timeout := defaultBashTimeout + if input.Timeout > 0 { + timeout = time.Duration(input.Timeout) * time.Millisecond + } + + execCtx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + + // Use sandbox if available, otherwise execute directly on host + if sb != nil { + return executeBashViaSandbox(execCtx, sb, input.Command, workdir, timeout) + } + + return executeBashDirect(execCtx, input.Command, workdir, timeout) + }, + ) +} + +// executeBashViaSandbox executes the command through the sandbox interface +func executeBashViaSandbox(ctx context.Context, sb sandbox.Sandbox, command, workdir string, timeout time.Duration) (BashOutput, error) { + // Prepend cd to workdir if different from sandbox workspace + fullCommand := command + if workdir != "" && workdir != sb.Config().WorkspacePath { + fullCommand = fmt.Sprintf("cd %q && %s", workdir, command) + } + + result, err := sb.Exec(ctx, "bash", "-c", fullCommand) + if err != nil { + return BashOutput{}, fmt.Errorf("sandbox exec failed: %w", err) + } + + output := result.Stdout + if result.Stderr != "" { + if output != "" { + output += "\n" + } + output += result.Stderr + } + + timedOut := result.Killed && strings.Contains(result.KillReason, "timeout") + + if len(output) > maxBashOutputLength { + output = output[:maxBashOutputLength] + output += fmt.Sprintf("\n\n\nbash tool truncated output as it exceeded %d char limit\n", maxBashOutputLength) + } + + if timedOut { + output += fmt.Sprintf("\n\n\nbash tool terminated command after exceeding timeout %v\n", timeout) + } + + return BashOutput{ + Output: output, + ExitCode: result.ExitCode, + TimedOut: timedOut, + }, nil +} + +// executeBashDirect executes the command directly on the host (no sandbox) +func executeBashDirect(ctx context.Context, command, workdir string, timeout time.Duration) (BashOutput, error) { + cmd := exec.CommandContext(ctx, "bash", "-c", command) + cmd.Dir = workdir + + var stdout, stderr bytes.Buffer + cmd.Stdout = &stdout + cmd.Stderr = &stderr + + err := cmd.Run() + + output := stdout.String() + if stderr.Len() > 0 { + if output != "" { + output += "\n" + } + output += stderr.String() + } + + timedOut := ctx.Err() == context.DeadlineExceeded + + if len(output) > maxBashOutputLength { + output = output[:maxBashOutputLength] + output += fmt.Sprintf("\n\n\nbash tool truncated output as it exceeded %d char limit\n", maxBashOutputLength) + } + + if timedOut { + output += fmt.Sprintf("\n\n\nbash tool terminated command after exceeding timeout %v\n", timeout) + } + + exitCode := 0 + if err != nil { + if exitErr, ok := err.(*exec.ExitError); ok { + exitCode = exitErr.ExitCode() + } else if !timedOut { + return BashOutput{}, fmt.Errorf("command execution failed: %w", err) + } + } + + return BashOutput{ + Output: output, + ExitCode: exitCode, + TimedOut: timedOut, + }, nil +} + +func isPathInWorkspace(path, workspacePath string) bool { + absPath, err := filepath.Abs(path) + if err != nil { + return false + } + absWorkspace, err := filepath.Abs(workspacePath) + if err != nil { + return false + } + return strings.HasPrefix(absPath, absWorkspace) +} diff --git a/pkg/harness/tools/edit.go b/pkg/harness/tools/edit.go new file mode 100644 index 00000000..e43dd05a --- /dev/null +++ b/pkg/harness/tools/edit.go @@ -0,0 +1,110 @@ +package tools + +import ( + "fmt" + "os" + "path/filepath" + "strings" + + "station/pkg/harness/sandbox" + + "github.com/firebase/genkit/go/ai" + "github.com/firebase/genkit/go/genkit" +) + +type EditInput struct { + FilePath string `json:"file_path" jsonschema:"description=Absolute or relative path to the file to edit"` + OldString string `json:"old_string" jsonschema:"description=The exact text to replace (must match exactly including whitespace)"` + NewString string `json:"new_string" jsonschema:"description=The text to replace it with"` + ReplaceAll bool `json:"replace_all,omitempty" jsonschema:"description=Replace all occurrences instead of just the first (default false)"` +} + +type EditOutput struct { + Success bool `json:"success"` + Path string `json:"path"` + Replacements int `json:"replacements"` +} + +func NewEditTool(genkitApp *genkit.Genkit, workspacePath string) ai.Tool { + return NewEditToolWithSandbox(genkitApp, workspacePath, nil) +} + +func NewEditToolWithSandbox(genkitApp *genkit.Genkit, workspacePath string, sb sandbox.Sandbox) ai.Tool { + return genkit.DefineTool( + genkitApp, + "edit", + `Perform exact string replacements in files. + +IMPORTANT: +- oldString must match EXACTLY (including indentation/whitespace) +- The edit will FAIL if oldString is not found +- The edit will FAIL if oldString is found multiple times (unless replaceAll is true) +- Use replaceAll=true for renaming variables across a file + +Examples: +- Fix a typo: oldString="teh", newString="the" +- Rename variable: oldString="oldName", newString="newName", replaceAll=true`, + func(ctx *ai.ToolContext, input EditInput) (EditOutput, error) { + if input.FilePath == "" { + return EditOutput{}, fmt.Errorf("file_path is required") + } + if input.OldString == "" { + return EditOutput{}, fmt.Errorf("old_string is required") + } + if input.OldString == input.NewString { + return EditOutput{}, fmt.Errorf("old_string and new_string must be different") + } + + path := input.FilePath + if !filepath.IsAbs(path) { + path = filepath.Join(workspacePath, path) + } + + var content []byte + var err error + + if sb != nil { + content, err = sb.ReadFile(ctx, path) + } else { + content, err = os.ReadFile(path) + } + if err != nil { + return EditOutput{}, fmt.Errorf("failed to read file: %w", err) + } + + fileContent := string(content) + count := strings.Count(fileContent, input.OldString) + + if count == 0 { + return EditOutput{}, fmt.Errorf("oldString not found in file content") + } + + if count > 1 && !input.ReplaceAll { + return EditOutput{}, fmt.Errorf("oldString found %d times - requires more context to uniquely identify or set replaceAll=true", count) + } + + var newContent string + if input.ReplaceAll { + newContent = strings.ReplaceAll(fileContent, input.OldString, input.NewString) + } else { + newContent = strings.Replace(fileContent, input.OldString, input.NewString, 1) + count = 1 + } + + if sb != nil { + err = sb.WriteFile(ctx, path, []byte(newContent), 0644) + } else { + err = os.WriteFile(path, []byte(newContent), 0644) + } + if err != nil { + return EditOutput{}, fmt.Errorf("failed to write file: %w", err) + } + + return EditOutput{ + Success: true, + Path: path, + Replacements: count, + }, nil + }, + ) +} diff --git a/pkg/harness/tools/git_tools.go b/pkg/harness/tools/git_tools.go new file mode 100644 index 00000000..1a472bda --- /dev/null +++ b/pkg/harness/tools/git_tools.go @@ -0,0 +1,514 @@ +package tools + +import ( + "context" + "fmt" + "os/exec" + "path/filepath" + "strings" + + "github.com/firebase/genkit/go/ai" + "github.com/firebase/genkit/go/genkit" +) + +type GitStatusInput struct { + Path string `json:"path,omitempty" jsonschema:"description=Subdirectory path relative to workspace (optional)"` +} + +type GitStatusOutput struct { + Branch string `json:"branch"` + Ahead int `json:"ahead"` + Behind int `json:"behind"` + Staged []string `json:"staged"` + Modified []string `json:"modified"` + Untracked []string `json:"untracked"` + Clean bool `json:"clean"` + Output string `json:"output"` +} + +func NewGitStatusTool(g *genkit.Genkit, workspacePath string) ai.Tool { + return genkit.DefineTool(g, "git_status", + "Get the current git status including branch, staged files, modified files, and untracked files", + func(ctx *ai.ToolContext, input GitStatusInput) (GitStatusOutput, error) { + repoPath := workspacePath + if input.Path != "" { + repoPath = filepath.Join(workspacePath, input.Path) + } + + execCtx, cancel := context.WithTimeout(ctx, defaultBashTimeout) + defer cancel() + + cmd := exec.CommandContext(execCtx, "git", "status", "--porcelain=v2", "--branch") + cmd.Dir = repoPath + output, err := cmd.Output() + if err != nil { + return GitStatusOutput{}, fmt.Errorf("git status failed: %w", err) + } + + result := GitStatusOutput{ + Output: string(output), + } + + lines := strings.Split(string(output), "\n") + for _, line := range lines { + if strings.HasPrefix(line, "# branch.head ") { + result.Branch = strings.TrimPrefix(line, "# branch.head ") + } else if strings.HasPrefix(line, "# branch.ab ") { + parts := strings.Fields(line) + if len(parts) >= 4 { + fmt.Sscanf(parts[2], "+%d", &result.Ahead) + fmt.Sscanf(parts[3], "-%d", &result.Behind) + } + } else if len(line) > 2 { + xy := line[:2] + path := "" + if strings.HasPrefix(line, "1 ") || strings.HasPrefix(line, "2 ") { + parts := strings.Fields(line) + if len(parts) >= 9 { + path = parts[8] + } + } else if strings.HasPrefix(line, "? ") { + path = strings.TrimPrefix(line, "? ") + result.Untracked = append(result.Untracked, path) + continue + } + + if path == "" { + continue + } + + if xy[0] != '.' && xy[0] != '?' { + result.Staged = append(result.Staged, path) + } + if xy[1] != '.' && xy[1] != '?' { + result.Modified = append(result.Modified, path) + } + } + } + + result.Clean = len(result.Staged) == 0 && len(result.Modified) == 0 && len(result.Untracked) == 0 + + return result, nil + }, + ) +} + +type GitDiffInput struct { + Path string `json:"path,omitempty" jsonschema:"description=Subdirectory path relative to workspace (optional)"` + Staged bool `json:"staged,omitempty" jsonschema:"description=Show staged changes instead of working directory changes"` + File string `json:"file,omitempty" jsonschema:"description=Specific file to diff"` +} + +type GitDiffOutput struct { + Diff string `json:"diff"` + FilesCount int `json:"files_count"` + Insertions int `json:"insertions"` + Deletions int `json:"deletions"` +} + +func NewGitDiffTool(g *genkit.Genkit, workspacePath string) ai.Tool { + return genkit.DefineTool(g, "git_diff", + "Show changes in the working directory or staged changes. Use staged=true to see staged changes.", + func(ctx *ai.ToolContext, input GitDiffInput) (GitDiffOutput, error) { + repoPath := workspacePath + if input.Path != "" { + repoPath = filepath.Join(workspacePath, input.Path) + } + + execCtx, cancel := context.WithTimeout(ctx, defaultBashTimeout) + defer cancel() + + args := []string{"diff"} + if input.Staged { + args = append(args, "--staged") + } + args = append(args, "--stat") + + if input.File != "" { + args = append(args, "--", input.File) + } + + cmd := exec.CommandContext(execCtx, "git", args...) + cmd.Dir = repoPath + stat, _ := cmd.Output() + + args = []string{"diff"} + if input.Staged { + args = append(args, "--staged") + } + if input.File != "" { + args = append(args, "--", input.File) + } + + cmd = exec.CommandContext(execCtx, "git", args...) + cmd.Dir = repoPath + diff, err := cmd.Output() + if err != nil { + return GitDiffOutput{}, fmt.Errorf("git diff failed: %w", err) + } + + result := GitDiffOutput{ + Diff: string(diff), + } + + statLines := strings.Split(string(stat), "\n") + for _, line := range statLines { + if strings.Contains(line, "files changed") || strings.Contains(line, "file changed") { + fmt.Sscanf(line, " %d file", &result.FilesCount) + if strings.Contains(line, "insertion") { + fmt.Sscanf(strings.Split(line, ",")[1], " %d insertion", &result.Insertions) + } + if strings.Contains(line, "deletion") { + parts := strings.Split(line, ",") + for _, p := range parts { + if strings.Contains(p, "deletion") { + fmt.Sscanf(p, " %d deletion", &result.Deletions) + } + } + } + } + } + + if len(result.Diff) > maxBashOutputLength { + result.Diff = result.Diff[:maxBashOutputLength] + "\n... (truncated)" + } + + return result, nil + }, + ) +} + +type GitLogInput struct { + Path string `json:"path,omitempty" jsonschema:"description=Subdirectory path relative to workspace (optional)"` + Count int `json:"count,omitempty" jsonschema:"description=Number of commits to show (default 10, max 50)"` + File string `json:"file,omitempty" jsonschema:"description=Show commits affecting this file"` +} + +type GitLogCommit struct { + Hash string `json:"hash"` + Author string `json:"author"` + Date string `json:"date"` + Subject string `json:"subject"` +} + +type GitLogOutput struct { + Commits []GitLogCommit `json:"commits"` +} + +func NewGitLogTool(g *genkit.Genkit, workspacePath string) ai.Tool { + return genkit.DefineTool(g, "git_log", + "Show recent commit history. Use count to limit results (default 10).", + func(ctx *ai.ToolContext, input GitLogInput) (GitLogOutput, error) { + repoPath := workspacePath + if input.Path != "" { + repoPath = filepath.Join(workspacePath, input.Path) + } + + execCtx, cancel := context.WithTimeout(ctx, defaultBashTimeout) + defer cancel() + + count := input.Count + if count <= 0 { + count = 10 + } + if count > 50 { + count = 50 + } + + format := "%H%n%an%n%ai%n%s%n---" + args := []string{"log", fmt.Sprintf("-n%d", count), fmt.Sprintf("--format=%s", format)} + + if input.File != "" { + args = append(args, "--", input.File) + } + + cmd := exec.CommandContext(execCtx, "git", args...) + cmd.Dir = repoPath + output, err := cmd.Output() + if err != nil { + return GitLogOutput{}, fmt.Errorf("git log failed: %w", err) + } + + result := GitLogOutput{} + entries := strings.Split(string(output), "---\n") + + for _, entry := range entries { + entry = strings.TrimSpace(entry) + if entry == "" { + continue + } + + lines := strings.Split(entry, "\n") + if len(lines) >= 4 { + result.Commits = append(result.Commits, GitLogCommit{ + Hash: lines[0], + Author: lines[1], + Date: lines[2], + Subject: lines[3], + }) + } + } + + return result, nil + }, + ) +} + +type GitAddInput struct { + Path string `json:"path,omitempty" jsonschema:"description=Subdirectory path relative to workspace (optional)"` + Files []string `json:"files,omitempty" jsonschema:"description=Specific files to stage"` + All bool `json:"all,omitempty" jsonschema:"description=Stage all changes including untracked files"` +} + +type GitAddOutput struct { + Added []string `json:"added"` + Message string `json:"message"` +} + +func NewGitAddTool(g *genkit.Genkit, workspacePath string) ai.Tool { + return genkit.DefineTool(g, "git_add", + "Stage files for commit. Use files array for specific files, or all=true to stage all changes.", + func(ctx *ai.ToolContext, input GitAddInput) (GitAddOutput, error) { + repoPath := workspacePath + if input.Path != "" { + repoPath = filepath.Join(workspacePath, input.Path) + } + + execCtx, cancel := context.WithTimeout(ctx, defaultBashTimeout) + defer cancel() + + args := []string{"add"} + if input.All { + args = append(args, "-A") + } else if len(input.Files) > 0 { + args = append(args, input.Files...) + } else { + return GitAddOutput{}, fmt.Errorf("must specify files or all=true") + } + + cmd := exec.CommandContext(execCtx, "git", args...) + cmd.Dir = repoPath + if err := cmd.Run(); err != nil { + return GitAddOutput{}, fmt.Errorf("git add failed: %w", err) + } + + result := GitAddOutput{ + Message: "Files staged successfully", + } + + if input.All { + result.Added = []string{"all changes"} + } else { + result.Added = input.Files + } + + return result, nil + }, + ) +} + +type GitCommitInput struct { + Path string `json:"path,omitempty" jsonschema:"description=Subdirectory path relative to workspace (optional)"` + Message string `json:"message" jsonschema:"description=Commit message (required)"` +} + +type GitCommitOutput struct { + Hash string `json:"hash"` + Message string `json:"message"` + Output string `json:"output"` +} + +func NewGitCommitTool(g *genkit.Genkit, workspacePath string) ai.Tool { + return genkit.DefineTool(g, "git_commit", + "Create a commit with staged changes. Requires a commit message.", + func(ctx *ai.ToolContext, input GitCommitInput) (GitCommitOutput, error) { + repoPath := workspacePath + if input.Path != "" { + repoPath = filepath.Join(workspacePath, input.Path) + } + + if input.Message == "" { + return GitCommitOutput{}, fmt.Errorf("commit message is required") + } + + execCtx, cancel := context.WithTimeout(ctx, defaultBashTimeout) + defer cancel() + + cmd := exec.CommandContext(execCtx, "git", "commit", "-m", input.Message) + cmd.Dir = repoPath + output, err := cmd.CombinedOutput() + if err != nil { + return GitCommitOutput{}, fmt.Errorf("git commit failed: %s", string(output)) + } + + cmd = exec.CommandContext(execCtx, "git", "rev-parse", "HEAD") + cmd.Dir = repoPath + hashBytes, _ := cmd.Output() + + return GitCommitOutput{ + Hash: strings.TrimSpace(string(hashBytes)), + Message: input.Message, + Output: string(output), + }, nil + }, + ) +} + +type GitBranchInput struct { + Path string `json:"path,omitempty" jsonschema:"description=Subdirectory path relative to workspace (optional)"` + Name string `json:"name,omitempty" jsonschema:"description=Branch name for create/checkout operations"` + Create bool `json:"create,omitempty" jsonschema:"description=Create a new branch with the given name"` + Checkout bool `json:"checkout,omitempty" jsonschema:"description=Switch to the branch with the given name"` + ListRemote bool `json:"list_remote,omitempty" jsonschema:"description=Include remote branches in listing"` +} + +type GitBranch struct { + Name string `json:"name"` + Current bool `json:"current"` + Remote bool `json:"remote"` +} + +type GitBranchOutput struct { + Branches []GitBranch `json:"branches"` + Current string `json:"current"` + Message string `json:"message,omitempty"` +} + +func NewGitBranchTool(g *genkit.Genkit, workspacePath string) ai.Tool { + return genkit.DefineTool(g, "git_branch", + "List, create, or switch branches. Use name+create to create new branch, name+checkout to switch.", + func(ctx *ai.ToolContext, input GitBranchInput) (GitBranchOutput, error) { + repoPath := workspacePath + if input.Path != "" { + repoPath = filepath.Join(workspacePath, input.Path) + } + + execCtx, cancel := context.WithTimeout(ctx, defaultBashTimeout) + defer cancel() + + if input.Name != "" && input.Create { + args := []string{"checkout", "-b", input.Name} + cmd := exec.CommandContext(execCtx, "git", args...) + cmd.Dir = repoPath + if err := cmd.Run(); err != nil { + return GitBranchOutput{}, fmt.Errorf("git branch create failed: %w", err) + } + return GitBranchOutput{ + Current: input.Name, + Message: fmt.Sprintf("Created and switched to branch '%s'", input.Name), + }, nil + } + + if input.Name != "" && input.Checkout { + cmd := exec.CommandContext(execCtx, "git", "checkout", input.Name) + cmd.Dir = repoPath + if err := cmd.Run(); err != nil { + return GitBranchOutput{}, fmt.Errorf("git checkout failed: %w", err) + } + return GitBranchOutput{ + Current: input.Name, + Message: fmt.Sprintf("Switched to branch '%s'", input.Name), + }, nil + } + + args := []string{"branch"} + if input.ListRemote { + args = append(args, "-a") + } + + cmd := exec.CommandContext(execCtx, "git", args...) + cmd.Dir = repoPath + output, err := cmd.Output() + if err != nil { + return GitBranchOutput{}, fmt.Errorf("git branch list failed: %w", err) + } + + result := GitBranchOutput{} + lines := strings.Split(string(output), "\n") + + for _, line := range lines { + line = strings.TrimSpace(line) + if line == "" { + continue + } + + branch := GitBranch{} + if strings.HasPrefix(line, "* ") { + branch.Current = true + branch.Name = strings.TrimPrefix(line, "* ") + result.Current = branch.Name + } else { + branch.Name = line + } + + if strings.HasPrefix(branch.Name, "remotes/") { + branch.Remote = true + branch.Name = strings.TrimPrefix(branch.Name, "remotes/") + } + + result.Branches = append(result.Branches, branch) + } + + return result, nil + }, + ) +} + +type GitResetInput struct { + Path string `json:"path,omitempty" jsonschema:"description=Subdirectory path relative to workspace (optional)"` + Files []string `json:"files,omitempty" jsonschema:"description=Specific files to unstage (empty = all)"` +} + +type GitResetOutput struct { + Reset []string `json:"reset"` + Message string `json:"message"` +} + +func NewGitResetTool(g *genkit.Genkit, workspacePath string) ai.Tool { + return genkit.DefineTool(g, "git_reset", + "Unstage files (soft reset). Specify files to unstage, or leave empty to unstage all.", + func(ctx *ai.ToolContext, input GitResetInput) (GitResetOutput, error) { + repoPath := workspacePath + if input.Path != "" { + repoPath = filepath.Join(workspacePath, input.Path) + } + + execCtx, cancel := context.WithTimeout(ctx, defaultBashTimeout) + defer cancel() + + args := []string{"reset", "HEAD"} + if len(input.Files) > 0 { + args = append(args, "--") + args = append(args, input.Files...) + } + + cmd := exec.CommandContext(execCtx, "git", args...) + cmd.Dir = repoPath + if err := cmd.Run(); err != nil { + return GitResetOutput{}, fmt.Errorf("git reset failed: %w", err) + } + + result := GitResetOutput{ + Message: "Files unstaged successfully", + } + + if len(input.Files) > 0 { + result.Reset = input.Files + } else { + result.Reset = []string{"all staged files"} + } + + return result, nil + }, + ) +} + +func (r *ToolRegistry) RegisterGitTools() { + r.tools["git_status"] = NewGitStatusTool(r.genkitApp, r.workspacePath) + r.tools["git_diff"] = NewGitDiffTool(r.genkitApp, r.workspacePath) + r.tools["git_log"] = NewGitLogTool(r.genkitApp, r.workspacePath) + r.tools["git_add"] = NewGitAddTool(r.genkitApp, r.workspacePath) + r.tools["git_commit"] = NewGitCommitTool(r.genkitApp, r.workspacePath) + r.tools["git_branch"] = NewGitBranchTool(r.genkitApp, r.workspacePath) + r.tools["git_reset"] = NewGitResetTool(r.genkitApp, r.workspacePath) +} diff --git a/pkg/harness/tools/glob.go b/pkg/harness/tools/glob.go new file mode 100644 index 00000000..2b60479c --- /dev/null +++ b/pkg/harness/tools/glob.go @@ -0,0 +1,202 @@ +package tools + +import ( + "fmt" + "os" + "path/filepath" + "sort" + "strings" + "time" + + "station/pkg/harness/sandbox" + + "github.com/firebase/genkit/go/ai" + "github.com/firebase/genkit/go/genkit" +) + +const ( + maxGlobResults = 100 + globTimeout = 60 * time.Second +) + +type GlobInput struct { + Pattern string `json:"pattern" jsonschema:"description=Glob pattern to match files (e.g., '**/*.go', 'src/**/*.ts')"` + Path string `json:"path,omitempty" jsonschema:"description=Base directory to search from (default: workspace root)"` +} + +type GlobOutput struct { + Files []string `json:"files"` + Count int `json:"count"` + Truncated bool `json:"truncated,omitempty"` +} + +func NewGlobTool(genkitApp *genkit.Genkit, workspacePath string) ai.Tool { + return NewGlobToolWithSandbox(genkitApp, workspacePath, nil) +} + +func NewGlobToolWithSandbox(genkitApp *genkit.Genkit, workspacePath string, sb sandbox.Sandbox) ai.Tool { + return genkit.DefineTool( + genkitApp, + "glob", + `Find files matching a glob pattern. + +Supports patterns like: +- "**/*.go" - all Go files recursively +- "src/**/*.ts" - TypeScript files under src/ +- "*.json" - JSON files in current directory +- "test_*.py" - Python test files + +Results are sorted by modification time (newest first). +Limited to 100 results with 60s timeout.`, + func(ctx *ai.ToolContext, input GlobInput) (GlobOutput, error) { + if input.Pattern == "" { + return GlobOutput{}, fmt.Errorf("pattern is required") + } + + basePath := workspacePath + if input.Path != "" { + if filepath.IsAbs(input.Path) { + basePath = input.Path + } else { + basePath = filepath.Join(workspacePath, input.Path) + } + } + + var matches []fileInfo + pattern := input.Pattern + hasDoublestar := strings.Contains(pattern, "**") + + if sb != nil { + matches = globViaSandbox(ctx, sb, basePath, pattern, hasDoublestar) + } else { + matches = globDirect(basePath, pattern, hasDoublestar) + } + + sort.Slice(matches, func(i, j int) bool { + return matches[i].modTime.After(matches[j].modTime) + }) + + truncated := false + if len(matches) > maxGlobResults { + matches = matches[:maxGlobResults] + truncated = true + } + + files := make([]string, len(matches)) + for i, m := range matches { + files[i] = m.path + } + + return GlobOutput{ + Files: files, + Count: len(files), + Truncated: truncated, + }, nil + }, + ) +} + +func globViaSandbox(ctx *ai.ToolContext, sb sandbox.Sandbox, basePath, pattern string, hasDoublestar bool) []fileInfo { + var matches []fileInfo + walkSandboxDir(ctx, sb, basePath, basePath, pattern, hasDoublestar, &matches) + return matches +} + +func walkSandboxDir(ctx *ai.ToolContext, sb sandbox.Sandbox, basePath, currentPath, pattern string, hasDoublestar bool, matches *[]fileInfo) { + files, err := sb.ListFiles(ctx, currentPath) + if err != nil { + return + } + + for _, f := range files { + fullPath := filepath.Join(currentPath, f.Name) + relPath, _ := filepath.Rel(basePath, fullPath) + + if f.IsDir { + walkSandboxDir(ctx, sb, basePath, fullPath, pattern, hasDoublestar, matches) + continue + } + + var matched bool + if hasDoublestar { + matched = matchDoublestar(pattern, relPath) + } else { + matched, _ = filepath.Match(pattern, filepath.Base(relPath)) + } + + if matched { + *matches = append(*matches, fileInfo{ + path: relPath, + modTime: f.ModTime, + }) + } + } +} + +func globDirect(basePath, pattern string, hasDoublestar bool) []fileInfo { + var matches []fileInfo + + filepath.Walk(basePath, func(path string, info os.FileInfo, err error) error { + if err != nil { + return nil + } + + if info.IsDir() { + return nil + } + + relPath, err := filepath.Rel(basePath, path) + if err != nil { + return nil + } + + var matched bool + if hasDoublestar { + matched = matchDoublestar(pattern, relPath) + } else { + matched, _ = filepath.Match(pattern, filepath.Base(relPath)) + } + + if matched { + matches = append(matches, fileInfo{ + path: relPath, + modTime: info.ModTime(), + }) + } + + return nil + }) + + return matches +} + +type fileInfo struct { + path string + modTime time.Time +} + +func matchDoublestar(pattern, path string) bool { + patternParts := strings.Split(pattern, "**") + if len(patternParts) == 1 { + matched, _ := filepath.Match(pattern, path) + return matched + } + + if len(patternParts) == 2 { + prefix := strings.TrimSuffix(patternParts[0], "/") + suffix := strings.TrimPrefix(patternParts[1], "/") + + if prefix != "" && !strings.HasPrefix(path, prefix) { + return false + } + + if suffix != "" { + matched, _ := filepath.Match(suffix, filepath.Base(path)) + return matched + } + + return true + } + + return false +} diff --git a/pkg/harness/tools/grep.go b/pkg/harness/tools/grep.go new file mode 100644 index 00000000..12777c48 --- /dev/null +++ b/pkg/harness/tools/grep.go @@ -0,0 +1,307 @@ +package tools + +import ( + "bufio" + "fmt" + "os" + "path/filepath" + "regexp" + "sort" + "strings" + "time" + + "station/pkg/harness/sandbox" + + "github.com/firebase/genkit/go/ai" + "github.com/firebase/genkit/go/genkit" +) + +const ( + maxGrepMatches = 100 + maxGrepFileSize = 10 * 1024 * 1024 + grepContextLines = 2 +) + +type GrepInput struct { + Pattern string `json:"pattern" jsonschema:"description=Regular expression pattern to search for"` + Path string `json:"path,omitempty" jsonschema:"description=Directory to search in (default: workspace root)"` + Include string `json:"include,omitempty" jsonschema:"description=File pattern to include (e.g., '*.go', '*.{ts,tsx}')"` +} + +type GrepMatch struct { + File string `json:"file"` + Line int `json:"line"` + Content string `json:"content"` +} + +type GrepOutput struct { + Matches []GrepMatch `json:"matches"` + Count int `json:"count"` + Truncated bool `json:"truncated,omitempty"` +} + +func NewGrepTool(genkitApp *genkit.Genkit, workspacePath string) ai.Tool { + return NewGrepToolWithSandbox(genkitApp, workspacePath, nil) +} + +func NewGrepToolWithSandbox(genkitApp *genkit.Genkit, workspacePath string, sb sandbox.Sandbox) ai.Tool { + return genkit.DefineTool( + genkitApp, + "grep", + `Search file contents using regular expressions. + +Supports full regex syntax: +- "log.*Error" - matches "logError", "log_Error", etc. +- "function\s+\w+" - matches function definitions +- "TODO|FIXME" - matches either word + +Filter files with the include parameter: +- "*.go" - Go files only +- "*.{ts,tsx}" - TypeScript files + +Returns file paths with matches sorted by modification time. +Limited to 100 matches.`, + func(ctx *ai.ToolContext, input GrepInput) (GrepOutput, error) { + if input.Pattern == "" { + return GrepOutput{}, fmt.Errorf("pattern is required") + } + + re, err := regexp.Compile(input.Pattern) + if err != nil { + return GrepOutput{}, fmt.Errorf("invalid regex pattern: %w", err) + } + + basePath := workspacePath + if input.Path != "" { + if filepath.IsAbs(input.Path) { + basePath = input.Path + } else { + basePath = filepath.Join(workspacePath, input.Path) + } + } + + var filesSearched []fileWithTime + + if sb != nil { + filesSearched = collectFilesViaSandbox(ctx, sb, basePath, input.Include) + } else { + filesSearched = collectFilesDirect(basePath, input.Include) + } + + sort.Slice(filesSearched, func(i, j int) bool { + return filesSearched[i].modTime.After(filesSearched[j].modTime) + }) + + var matches []GrepMatch + for _, f := range filesSearched { + if len(matches) >= maxGrepMatches { + break + } + + var fileMatches []GrepMatch + if sb != nil { + fileMatches = searchFileViaSandbox(ctx, sb, f.path, re, basePath) + } else { + fileMatches, _ = searchFile(f.path, re, basePath) + } + + for _, m := range fileMatches { + if len(matches) >= maxGrepMatches { + break + } + matches = append(matches, m) + } + } + + truncated := len(matches) >= maxGrepMatches + + return GrepOutput{ + Matches: matches, + Count: len(matches), + Truncated: truncated, + }, nil + }, + ) +} + +func collectFilesViaSandbox(ctx *ai.ToolContext, sb sandbox.Sandbox, basePath, include string) []fileWithTime { + var files []fileWithTime + walkSandboxDirForGrep(ctx, sb, basePath, basePath, include, &files) + return files +} + +func walkSandboxDirForGrep(ctx *ai.ToolContext, sb sandbox.Sandbox, basePath, currentPath, include string, files *[]fileWithTime) { + entries, err := sb.ListFiles(ctx, currentPath) + if err != nil { + return + } + + for _, entry := range entries { + if entry.IsDir { + if shouldSkipDir(entry.Name) { + continue + } + walkSandboxDirForGrep(ctx, sb, basePath, filepath.Join(currentPath, entry.Name), include, files) + continue + } + + if entry.Size > maxGrepFileSize { + continue + } + + if include != "" && !matchInclude(include, entry.Name) { + continue + } + + *files = append(*files, fileWithTime{ + path: filepath.Join(currentPath, entry.Name), + modTime: entry.ModTime, + }) + } +} + +func searchFileViaSandbox(ctx *ai.ToolContext, sb sandbox.Sandbox, path string, re *regexp.Regexp, basePath string) []GrepMatch { + content, err := sb.ReadFile(ctx, path) + if err != nil { + return nil + } + + relPath, err := filepath.Rel(basePath, path) + if err != nil { + relPath = path + } + + var matches []GrepMatch + lines := strings.Split(string(content), "\n") + + for lineNum, line := range lines { + if re.MatchString(line) { + matches = append(matches, GrepMatch{ + File: relPath, + Line: lineNum + 1, + Content: truncateLine(line, 200), + }) + } + } + + return matches +} + +func collectFilesDirect(basePath, include string) []fileWithTime { + var files []fileWithTime + + filepath.Walk(basePath, func(path string, info os.FileInfo, err error) error { + if err != nil { + return nil + } + + if info.IsDir() { + if shouldSkipDir(info.Name()) { + return filepath.SkipDir + } + return nil + } + + if info.Size() > maxGrepFileSize { + return nil + } + + if include != "" && !matchInclude(include, info.Name()) { + return nil + } + + files = append(files, fileWithTime{ + path: path, + modTime: info.ModTime(), + }) + + return nil + }) + + return files +} + +type fileWithTime struct { + path string + modTime time.Time +} + +func searchFile(path string, re *regexp.Regexp, basePath string) ([]GrepMatch, error) { + file, err := os.Open(path) + if err != nil { + return nil, err + } + defer file.Close() + + relPath, err := filepath.Rel(basePath, path) + if err != nil { + relPath = path + } + + var matches []GrepMatch + scanner := bufio.NewScanner(file) + lineNum := 0 + + for scanner.Scan() { + lineNum++ + line := scanner.Text() + + if re.MatchString(line) { + matches = append(matches, GrepMatch{ + File: relPath, + Line: lineNum, + Content: truncateLine(line, 200), + }) + } + } + + return matches, scanner.Err() +} + +func shouldSkipDir(name string) bool { + skipDirs := map[string]bool{ + ".git": true, + "node_modules": true, + "vendor": true, + "__pycache__": true, + ".venv": true, + "venv": true, + ".idea": true, + ".vscode": true, + "dist": true, + "build": true, + ".next": true, + "coverage": true, + } + return skipDirs[name] +} + +func matchInclude(pattern, filename string) bool { + if strings.Contains(pattern, "{") { + start := strings.Index(pattern, "{") + end := strings.Index(pattern, "}") + if start != -1 && end != -1 && end > start { + prefix := pattern[:start] + suffix := pattern[end+1:] + extensions := strings.Split(pattern[start+1:end], ",") + + for _, ext := range extensions { + fullPattern := prefix + strings.TrimSpace(ext) + suffix + if matched, _ := filepath.Match(fullPattern, filename); matched { + return true + } + } + return false + } + } + + matched, _ := filepath.Match(pattern, filename) + return matched +} + +func truncateLine(line string, maxLen int) string { + if len(line) <= maxLen { + return line + } + return line[:maxLen] + "..." +} diff --git a/pkg/harness/tools/read.go b/pkg/harness/tools/read.go new file mode 100644 index 00000000..cb4751d9 --- /dev/null +++ b/pkg/harness/tools/read.go @@ -0,0 +1,112 @@ +package tools + +import ( + "fmt" + "os" + "path/filepath" + "strings" + + "station/pkg/harness/sandbox" + + "github.com/firebase/genkit/go/ai" + "github.com/firebase/genkit/go/genkit" +) + +const ( + defaultReadLimit = 2000 + maxLineLength = 2000 +) + +type ReadInput struct { + FilePath string `json:"file_path" jsonschema:"description=Absolute or relative path to the file to read"` + Offset int `json:"offset,omitempty" jsonschema:"description=Line number to start reading from (0-based, default 0)"` + Limit int `json:"limit,omitempty" jsonschema:"description=Maximum number of lines to read (default 2000)"` +} + +type ReadOutput struct { + Content string `json:"content"` + TotalLines int `json:"total_lines"` + StartLine int `json:"start_line"` + EndLine int `json:"end_line"` +} + +func NewReadTool(genkitApp *genkit.Genkit, workspacePath string) ai.Tool { + return NewReadToolWithSandbox(genkitApp, workspacePath, nil) +} + +func NewReadToolWithSandbox(genkitApp *genkit.Genkit, workspacePath string, sb sandbox.Sandbox) ai.Tool { + return genkit.DefineTool( + genkitApp, + "read", + `Read a file from the filesystem. + +Returns file content with line numbers (cat -n format). +By default reads up to 2000 lines from the beginning. +Use offset and limit for reading specific sections of large files. +Lines longer than 2000 characters are truncated.`, + func(ctx *ai.ToolContext, input ReadInput) (ReadOutput, error) { + if input.FilePath == "" { + return ReadOutput{}, fmt.Errorf("file_path is required") + } + + path := input.FilePath + if !filepath.IsAbs(path) { + path = filepath.Join(workspacePath, path) + } + + var content []byte + var err error + + if sb != nil { + content, err = sb.ReadFile(ctx, path) + } else { + content, err = os.ReadFile(path) + } + if err != nil { + return ReadOutput{}, fmt.Errorf("failed to read file: %w", err) + } + + lines := strings.Split(string(content), "\n") + totalLines := len(lines) + + offset := input.Offset + if offset < 0 { + offset = 0 + } + if offset >= totalLines { + return ReadOutput{ + Content: "", + TotalLines: totalLines, + StartLine: offset, + EndLine: offset, + }, nil + } + + limit := input.Limit + if limit <= 0 { + limit = defaultReadLimit + } + + endLine := offset + limit + if endLine > totalLines { + endLine = totalLines + } + + var result strings.Builder + for i := offset; i < endLine; i++ { + line := lines[i] + if len(line) > maxLineLength { + line = line[:maxLineLength] + "..." + } + result.WriteString(fmt.Sprintf("%6d\t%s\n", i+1, line)) + } + + return ReadOutput{ + Content: result.String(), + TotalLines: totalLines, + StartLine: offset + 1, + EndLine: endLine, + }, nil + }, + ) +} diff --git a/pkg/harness/tools/registry.go b/pkg/harness/tools/registry.go new file mode 100644 index 00000000..d53ac852 --- /dev/null +++ b/pkg/harness/tools/registry.go @@ -0,0 +1,91 @@ +package tools + +import ( + "station/pkg/harness/sandbox" + + "github.com/firebase/genkit/go/ai" + "github.com/firebase/genkit/go/genkit" +) + +type ToolRegistry struct { + genkitApp *genkit.Genkit + workspacePath string + sandbox sandbox.Sandbox + tools map[string]ai.Tool +} + +func NewToolRegistry(genkitApp *genkit.Genkit, workspacePath string) *ToolRegistry { + return &ToolRegistry{ + genkitApp: genkitApp, + workspacePath: workspacePath, + tools: make(map[string]ai.Tool), + } +} + +func NewToolRegistryWithSandbox(genkitApp *genkit.Genkit, workspacePath string, sb sandbox.Sandbox) *ToolRegistry { + return &ToolRegistry{ + genkitApp: genkitApp, + workspacePath: workspacePath, + sandbox: sb, + tools: make(map[string]ai.Tool), + } +} + +func (r *ToolRegistry) RegisterBuiltinTools() error { + bashTool := NewBashToolWithSandbox(r.genkitApp, r.workspacePath, r.sandbox) + r.tools["bash"] = bashTool + + readTool := NewReadToolWithSandbox(r.genkitApp, r.workspacePath, r.sandbox) + r.tools["read"] = readTool + + writeTool := NewWriteToolWithSandbox(r.genkitApp, r.workspacePath, r.sandbox) + r.tools["write"] = writeTool + + editTool := NewEditToolWithSandbox(r.genkitApp, r.workspacePath, r.sandbox) + r.tools["edit"] = editTool + + globTool := NewGlobToolWithSandbox(r.genkitApp, r.workspacePath, r.sandbox) + r.tools["glob"] = globTool + + grepTool := NewGrepToolWithSandbox(r.genkitApp, r.workspacePath, r.sandbox) + r.tools["grep"] = grepTool + + return nil +} + +func (r *ToolRegistry) Sandbox() sandbox.Sandbox { + return r.sandbox +} + +func (r *ToolRegistry) Get(name string) (ai.Tool, bool) { + tool, ok := r.tools[name] + return tool, ok +} + +func (r *ToolRegistry) All() []ai.Tool { + result := make([]ai.Tool, 0, len(r.tools)) + for _, tool := range r.tools { + result = append(result, tool) + } + return result +} + +func (r *ToolRegistry) AllRefs() []ai.ToolRef { + result := make([]ai.ToolRef, 0, len(r.tools)) + for _, tool := range r.tools { + result = append(result, tool) + } + return result +} + +func (r *ToolRegistry) Names() []string { + result := make([]string, 0, len(r.tools)) + for name := range r.tools { + result = append(result, name) + } + return result +} + +func (r *ToolRegistry) Register(name string, tool ai.Tool) { + r.tools[name] = tool +} diff --git a/pkg/harness/tools/write.go b/pkg/harness/tools/write.go new file mode 100644 index 00000000..630964a9 --- /dev/null +++ b/pkg/harness/tools/write.go @@ -0,0 +1,74 @@ +package tools + +import ( + "fmt" + "os" + "path/filepath" + + "station/pkg/harness/sandbox" + + "github.com/firebase/genkit/go/ai" + "github.com/firebase/genkit/go/genkit" +) + +type WriteInput struct { + FilePath string `json:"file_path" jsonschema:"description=Absolute or relative path to the file to write"` + Content string `json:"content" jsonschema:"description=Content to write to the file"` +} + +type WriteOutput struct { + Success bool `json:"success"` + Path string `json:"path"` + Bytes int `json:"bytes"` +} + +func NewWriteTool(genkitApp *genkit.Genkit, workspacePath string) ai.Tool { + return NewWriteToolWithSandbox(genkitApp, workspacePath, nil) +} + +func NewWriteToolWithSandbox(genkitApp *genkit.Genkit, workspacePath string, sb sandbox.Sandbox) ai.Tool { + return genkit.DefineTool( + genkitApp, + "write", + `Write content to a file, creating it if it doesn't exist or overwriting if it does. + +IMPORTANT: +- Creates parent directories automatically if they don't exist +- Overwrites existing files without warning +- Use edit tool for making changes to existing files`, + func(ctx *ai.ToolContext, input WriteInput) (WriteOutput, error) { + if input.FilePath == "" { + return WriteOutput{}, fmt.Errorf("file_path is required") + } + + path := input.FilePath + if !filepath.IsAbs(path) { + path = filepath.Join(workspacePath, path) + } + + if sb != nil { + dir := filepath.Dir(path) + sb.Exec(ctx, "mkdir", "-p", dir) + + if err := sb.WriteFile(ctx, path, []byte(input.Content), 0644); err != nil { + return WriteOutput{}, fmt.Errorf("failed to write file: %w", err) + } + } else { + dir := filepath.Dir(path) + if err := os.MkdirAll(dir, 0755); err != nil { + return WriteOutput{}, fmt.Errorf("failed to create directory: %w", err) + } + + if err := os.WriteFile(path, []byte(input.Content), 0644); err != nil { + return WriteOutput{}, fmt.Errorf("failed to write file: %w", err) + } + } + + return WriteOutput{ + Success: true, + Path: path, + Bytes: len(input.Content), + }, nil + }, + ) +} diff --git a/pkg/harness/trace/tracer.go b/pkg/harness/trace/tracer.go new file mode 100644 index 00000000..1ec2c896 --- /dev/null +++ b/pkg/harness/trace/tracer.go @@ -0,0 +1,357 @@ +package trace + +import ( + "context" + "time" + + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/codes" + "go.opentelemetry.io/otel/trace" +) + +const TracerName = "station/pkg/harness" + +type Tracer struct { + tracer trace.Tracer +} + +func New() *Tracer { + return &Tracer{ + tracer: otel.Tracer(TracerName), + } +} + +func NewWithTracer(t trace.Tracer) *Tracer { + return &Tracer{tracer: t} +} + +type ExecutionSpan struct { + span trace.Span + startTime time.Time +} + +func (t *Tracer) StartExecution(ctx context.Context, runID, agentID, agentName string) (context.Context, *ExecutionSpan) { + ctx, span := t.tracer.Start(ctx, "agent_execution", + trace.WithAttributes( + attribute.String("harness.run_id", runID), + attribute.String("harness.agent_id", agentID), + attribute.String("harness.agent_name", agentName), + attribute.String("harness.type", "agentic"), + ), + ) + + return ctx, &ExecutionSpan{ + span: span, + startTime: time.Now(), + } +} + +func (s *ExecutionSpan) SetWorkflow(workflowID, workflowRunID, stepName string) { + s.span.SetAttributes( + attribute.String("harness.workflow_id", workflowID), + attribute.String("harness.workflow_run_id", workflowRunID), + attribute.String("harness.step_name", stepName), + ) +} + +func (s *ExecutionSpan) End(status string, totalSteps int, err error) { + duration := time.Since(s.startTime) + + s.span.SetAttributes( + attribute.String("harness.status", status), + attribute.Int("harness.total_steps", totalSteps), + attribute.Int64("harness.duration_ms", duration.Milliseconds()), + ) + + if err != nil { + s.span.RecordError(err) + s.span.SetStatus(codes.Error, err.Error()) + } else { + s.span.SetStatus(codes.Ok, "") + } + + s.span.End() +} + +type SetupSpan struct { + span trace.Span + startTime time.Time +} + +func (t *Tracer) StartSetup(ctx context.Context) (context.Context, *SetupSpan) { + ctx, span := t.tracer.Start(ctx, "harness_setup") + return ctx, &SetupSpan{ + span: span, + startTime: time.Now(), + } +} + +func (s *SetupSpan) SetWorkspace(path, mode string) { + s.span.SetAttributes( + attribute.String("setup.workspace_path", path), + attribute.String("setup.workspace_mode", mode), + ) +} + +func (s *SetupSpan) SetGit(branch string, branchCreated bool) { + s.span.SetAttributes( + attribute.String("setup.git_branch", branch), + attribute.Bool("setup.branch_created", branchCreated), + ) +} + +func (s *SetupSpan) SetPreloadedFiles(count int, bytes int64) { + s.span.SetAttributes( + attribute.Int("setup.preload_files_count", count), + attribute.Int64("setup.preload_bytes", bytes), + ) +} + +func (s *SetupSpan) End(err error) { + duration := time.Since(s.startTime) + s.span.SetAttributes(attribute.Int64("setup.duration_ms", duration.Milliseconds())) + + if err != nil { + s.span.RecordError(err) + s.span.SetStatus(codes.Error, err.Error()) + } + s.span.End() +} + +type LoopStepSpan struct { + span trace.Span + startTime time.Time + stepNum int +} + +func (t *Tracer) StartLoopStep(ctx context.Context, stepNum int) (context.Context, *LoopStepSpan) { + ctx, span := t.tracer.Start(ctx, "agentic_loop_step", + trace.WithAttributes( + attribute.Int("step.number", stepNum), + ), + ) + + return ctx, &LoopStepSpan{ + span: span, + startTime: time.Now(), + stepNum: stepNum, + } +} + +func (s *LoopStepSpan) End(finishReason string) { + duration := time.Since(s.startTime) + s.span.SetAttributes( + attribute.Int64("step.duration_ms", duration.Milliseconds()), + attribute.String("step.finish_reason", finishReason), + ) + s.span.End() +} + +type LLMGenerateSpan struct { + span trace.Span + startTime time.Time +} + +func (t *Tracer) StartLLMGenerate(ctx context.Context, model string) (context.Context, *LLMGenerateSpan) { + ctx, span := t.tracer.Start(ctx, "llm_generate", + trace.WithAttributes( + attribute.String("llm.model", model), + ), + ) + + return ctx, &LLMGenerateSpan{ + span: span, + startTime: time.Now(), + } +} + +func (s *LLMGenerateSpan) SetTokens(input, output int) { + s.span.SetAttributes( + attribute.Int("llm.input_tokens", input), + attribute.Int("llm.output_tokens", output), + ) +} + +func (s *LLMGenerateSpan) SetToolRequests(tools []string) { + s.span.SetAttributes( + attribute.StringSlice("llm.tool_requests", tools), + attribute.Int("llm.tool_request_count", len(tools)), + ) +} + +func (s *LLMGenerateSpan) SetThinking(thinking string) { + if len(thinking) > 500 { + thinking = thinking[:500] + "..." + } + s.span.SetAttributes(attribute.String("llm.thinking", thinking)) +} + +func (s *LLMGenerateSpan) End(err error) { + duration := time.Since(s.startTime) + s.span.SetAttributes(attribute.Int64("llm.duration_ms", duration.Milliseconds())) + + if err != nil { + s.span.RecordError(err) + s.span.SetStatus(codes.Error, err.Error()) + } + s.span.End() +} + +type ToolExecutionSpan struct { + span trace.Span + startTime time.Time +} + +func (t *Tracer) StartToolExecution(ctx context.Context, toolName, callID string) (context.Context, *ToolExecutionSpan) { + ctx, span := t.tracer.Start(ctx, "tool_execution", + trace.WithAttributes( + attribute.String("tool.name", toolName), + attribute.String("tool.call_id", callID), + ), + ) + + return ctx, &ToolExecutionSpan{ + span: span, + startTime: time.Now(), + } +} + +func (s *ToolExecutionSpan) SetInput(input map[string]interface{}) { + for k, v := range input { + if str, ok := v.(string); ok { + if len(str) > 200 { + str = str[:200] + "..." + } + s.span.SetAttributes(attribute.String("tool.input."+k, str)) + } + } +} + +func (s *ToolExecutionSpan) SetOutput(outputBytes int, truncated bool) { + s.span.SetAttributes( + attribute.Int("tool.output_bytes", outputBytes), + attribute.Bool("tool.output_truncated", truncated), + ) +} + +func (s *ToolExecutionSpan) SetHookResult(hookType, result string) { + s.span.SetAttributes(attribute.String("tool.hook."+hookType, result)) +} + +func (s *ToolExecutionSpan) End(success bool, err error) { + duration := time.Since(s.startTime) + s.span.SetAttributes( + attribute.Int64("tool.duration_ms", duration.Milliseconds()), + attribute.Bool("tool.success", success), + ) + + if err != nil { + s.span.RecordError(err) + s.span.SetStatus(codes.Error, err.Error()) + } + s.span.End() +} + +type CompactionSpan struct { + span trace.Span + startTime time.Time +} + +func (t *Tracer) StartCompaction(ctx context.Context, currentTokens, threshold int) (context.Context, *CompactionSpan) { + ctx, span := t.tracer.Start(ctx, "compaction", + trace.WithAttributes( + attribute.Int("compaction.current_tokens", currentTokens), + attribute.Int("compaction.threshold", threshold), + ), + ) + + return ctx, &CompactionSpan{ + span: span, + startTime: time.Now(), + } +} + +func (s *CompactionSpan) SetResult(triggered bool, messagesBefore, messagesAfter, tokensBefore, tokensAfter int) { + s.span.SetAttributes( + attribute.Bool("compaction.triggered", triggered), + attribute.Int("compaction.messages_before", messagesBefore), + attribute.Int("compaction.messages_after", messagesAfter), + attribute.Int("compaction.tokens_before", tokensBefore), + attribute.Int("compaction.tokens_after", tokensAfter), + ) +} + +func (s *CompactionSpan) End(err error) { + duration := time.Since(s.startTime) + s.span.SetAttributes(attribute.Int64("compaction.duration_ms", duration.Milliseconds())) + + if err != nil { + s.span.RecordError(err) + s.span.SetStatus(codes.Error, err.Error()) + } + s.span.End() +} + +type CleanupSpan struct { + span trace.Span + startTime time.Time +} + +func (t *Tracer) StartCleanup(ctx context.Context) (context.Context, *CleanupSpan) { + ctx, span := t.tracer.Start(ctx, "harness_cleanup") + return ctx, &CleanupSpan{ + span: span, + startTime: time.Now(), + } +} + +func (s *CleanupSpan) SetGitCommit(sha string) { + s.span.SetAttributes(attribute.String("cleanup.git_commit", sha)) +} + +func (s *CleanupSpan) SetUploadedFiles(count int) { + s.span.SetAttributes(attribute.Int("cleanup.uploaded_files", count)) +} + +func (s *CleanupSpan) SetSavedContext(key string) { + s.span.SetAttributes(attribute.String("cleanup.context_key", key)) +} + +func (s *CleanupSpan) End(err error) { + duration := time.Since(s.startTime) + s.span.SetAttributes(attribute.Int64("cleanup.duration_ms", duration.Milliseconds())) + + if err != nil { + s.span.RecordError(err) + s.span.SetStatus(codes.Error, err.Error()) + } + s.span.End() +} + +func (t *Tracer) RecordEvent(ctx context.Context, name string, attrs ...attribute.KeyValue) { + span := trace.SpanFromContext(ctx) + span.AddEvent(name, trace.WithAttributes(attrs...)) +} + +func (t *Tracer) RecordDoomLoopDetected(ctx context.Context, toolName string, consecutiveCount int) { + t.RecordEvent(ctx, "doom_loop_detected", + attribute.String("doom_loop.tool", toolName), + attribute.Int("doom_loop.consecutive_count", consecutiveCount), + ) +} + +func (t *Tracer) RecordPermissionDenied(ctx context.Context, toolName, pattern, command string) { + t.RecordEvent(ctx, "permission_denied", + attribute.String("permission.tool", toolName), + attribute.String("permission.pattern", pattern), + attribute.String("permission.command", command), + ) +} + +func (t *Tracer) RecordApprovalRequired(ctx context.Context, toolName, reason string) { + t.RecordEvent(ctx, "approval_required", + attribute.String("approval.tool", toolName), + attribute.String("approval.reason", reason), + ) +} diff --git a/pkg/harness/workspace/host.go b/pkg/harness/workspace/host.go new file mode 100644 index 00000000..997cb755 --- /dev/null +++ b/pkg/harness/workspace/host.go @@ -0,0 +1,248 @@ +package workspace + +import ( + "context" + "fmt" + "os" + "path/filepath" + "strings" +) + +type HostWorkspace struct { + basePath string + allowedPaths []string + deniedPatterns []string + initialized bool +} + +type HostWorkspaceOption func(*HostWorkspace) + +func WithAllowedPaths(paths []string) HostWorkspaceOption { + return func(w *HostWorkspace) { + w.allowedPaths = paths + } +} + +func WithDeniedPatterns(patterns []string) HostWorkspaceOption { + return func(w *HostWorkspace) { + w.deniedPatterns = patterns + } +} + +func NewHostWorkspace(basePath string, opts ...HostWorkspaceOption) *HostWorkspace { + w := &HostWorkspace{ + basePath: basePath, + deniedPatterns: []string{ + ".git/objects/*", + ".git/hooks/*", + "node_modules/*", + "*.env", + "*.pem", + "*.key", + "*credentials*", + "*secrets*", + }, + } + + for _, opt := range opts { + opt(w) + } + + return w +} + +func (w *HostWorkspace) Initialize(ctx context.Context) error { + absPath, err := filepath.Abs(w.basePath) + if err != nil { + return fmt.Errorf("failed to resolve workspace path: %w", err) + } + w.basePath = absPath + + if _, err := os.Stat(w.basePath); os.IsNotExist(err) { + if err := os.MkdirAll(w.basePath, 0755); err != nil { + return fmt.Errorf("failed to create workspace directory: %w", err) + } + } + + w.initialized = true + return nil +} + +func (w *HostWorkspace) Path() string { + return w.basePath +} + +func (w *HostWorkspace) ReadFile(ctx context.Context, path string) ([]byte, error) { + if !w.initialized { + return nil, fmt.Errorf("workspace not initialized") + } + + fullPath, err := w.resolvePath(path) + if err != nil { + return nil, err + } + + if err := w.checkAccess(fullPath); err != nil { + return nil, err + } + + return os.ReadFile(fullPath) +} + +func (w *HostWorkspace) WriteFile(ctx context.Context, path string, data []byte) error { + if !w.initialized { + return fmt.Errorf("workspace not initialized") + } + + fullPath, err := w.resolvePath(path) + if err != nil { + return err + } + + if err := w.checkAccess(fullPath); err != nil { + return err + } + + dir := filepath.Dir(fullPath) + if err := os.MkdirAll(dir, 0755); err != nil { + return fmt.Errorf("failed to create directory: %w", err) + } + + return os.WriteFile(fullPath, data, 0644) +} + +func (w *HostWorkspace) DeleteFile(ctx context.Context, path string) error { + if !w.initialized { + return fmt.Errorf("workspace not initialized") + } + + fullPath, err := w.resolvePath(path) + if err != nil { + return err + } + + if err := w.checkAccess(fullPath); err != nil { + return err + } + + return os.Remove(fullPath) +} + +func (w *HostWorkspace) ListFiles(ctx context.Context, dir string) ([]string, error) { + if !w.initialized { + return nil, fmt.Errorf("workspace not initialized") + } + + fullPath, err := w.resolvePath(dir) + if err != nil { + return nil, err + } + + if err := w.checkAccess(fullPath); err != nil { + return nil, err + } + + entries, err := os.ReadDir(fullPath) + if err != nil { + return nil, err + } + + var files []string + for _, entry := range entries { + files = append(files, entry.Name()) + } + + return files, nil +} + +func (w *HostWorkspace) Exists(ctx context.Context, path string) (bool, error) { + if !w.initialized { + return false, fmt.Errorf("workspace not initialized") + } + + fullPath, err := w.resolvePath(path) + if err != nil { + return false, err + } + + _, err = os.Stat(fullPath) + if os.IsNotExist(err) { + return false, nil + } + if err != nil { + return false, err + } + + return true, nil +} + +func (w *HostWorkspace) Close(ctx context.Context) error { + w.initialized = false + return nil +} + +func (w *HostWorkspace) resolvePath(path string) (string, error) { + if filepath.IsAbs(path) { + return path, nil + } + + fullPath := filepath.Join(w.basePath, path) + fullPath, err := filepath.Abs(fullPath) + if err != nil { + return "", fmt.Errorf("failed to resolve path: %w", err) + } + + if !strings.HasPrefix(fullPath, w.basePath) { + return "", fmt.Errorf("path traversal detected: %s is outside workspace", path) + } + + return fullPath, nil +} + +func (w *HostWorkspace) checkAccess(path string) error { + relPath, err := filepath.Rel(w.basePath, path) + if err != nil { + return fmt.Errorf("failed to get relative path: %w", err) + } + + for _, pattern := range w.deniedPatterns { + matched, _ := filepath.Match(pattern, relPath) + if matched { + return fmt.Errorf("access denied: path matches denied pattern %s", pattern) + } + + if strings.Contains(pattern, "*") { + dir := filepath.Dir(relPath) + for dir != "." && dir != "/" { + matched, _ = filepath.Match(pattern, dir) + if matched { + return fmt.Errorf("access denied: path is under denied directory pattern %s", pattern) + } + dir = filepath.Dir(dir) + } + } + } + + if len(w.allowedPaths) > 0 { + allowed := false + for _, allowedPath := range w.allowedPaths { + if strings.HasPrefix(relPath, allowedPath) || relPath == allowedPath { + allowed = true + break + } + } + if !allowed { + return fmt.Errorf("access denied: path not in allowed paths") + } + } + + return nil +} + +func (w *HostWorkspace) IsInWorkspace(path string) bool { + fullPath, err := w.resolvePath(path) + if err != nil { + return false + } + return strings.HasPrefix(fullPath, w.basePath) +} diff --git a/pkg/harness/workspace/host_test.go b/pkg/harness/workspace/host_test.go new file mode 100644 index 00000000..3950cc42 --- /dev/null +++ b/pkg/harness/workspace/host_test.go @@ -0,0 +1,211 @@ +package workspace + +import ( + "context" + "os" + "path/filepath" + "testing" +) + +func TestHostWorkspace_Initialize(t *testing.T) { + tmpDir := t.TempDir() + + w := NewHostWorkspace(tmpDir) + ctx := context.Background() + + if err := w.Initialize(ctx); err != nil { + t.Fatalf("Initialize failed: %v", err) + } + + if !w.initialized { + t.Error("workspace should be initialized") + } + + if w.Path() != tmpDir { + t.Errorf("Path() = %q, want %q", w.Path(), tmpDir) + } +} + +func TestHostWorkspace_ReadWriteFile(t *testing.T) { + tmpDir := t.TempDir() + + w := NewHostWorkspace(tmpDir) + ctx := context.Background() + + if err := w.Initialize(ctx); err != nil { + t.Fatalf("Initialize failed: %v", err) + } + + content := []byte("hello world") + + if err := w.WriteFile(ctx, "test.txt", content); err != nil { + t.Fatalf("WriteFile failed: %v", err) + } + + read, err := w.ReadFile(ctx, "test.txt") + if err != nil { + t.Fatalf("ReadFile failed: %v", err) + } + + if string(read) != string(content) { + t.Errorf("ReadFile returned %q, want %q", string(read), string(content)) + } +} + +func TestHostWorkspace_NestedDirectory(t *testing.T) { + tmpDir := t.TempDir() + + w := NewHostWorkspace(tmpDir) + ctx := context.Background() + + if err := w.Initialize(ctx); err != nil { + t.Fatalf("Initialize failed: %v", err) + } + + content := []byte("nested content") + + if err := w.WriteFile(ctx, "a/b/c/test.txt", content); err != nil { + t.Fatalf("WriteFile failed: %v", err) + } + + if _, err := os.Stat(filepath.Join(tmpDir, "a/b/c/test.txt")); os.IsNotExist(err) { + t.Error("nested file was not created") + } +} + +func TestHostWorkspace_PathTraversal(t *testing.T) { + tmpDir := t.TempDir() + + w := NewHostWorkspace(tmpDir) + ctx := context.Background() + + if err := w.Initialize(ctx); err != nil { + t.Fatalf("Initialize failed: %v", err) + } + + _, err := w.ReadFile(ctx, "../../../etc/passwd") + if err == nil { + t.Error("expected error for path traversal") + } +} + +func TestHostWorkspace_DeniedPatterns(t *testing.T) { + tmpDir := t.TempDir() + + w := NewHostWorkspace(tmpDir, WithDeniedPatterns([]string{"*.secret", "config/*.env"})) + ctx := context.Background() + + if err := w.Initialize(ctx); err != nil { + t.Fatalf("Initialize failed: %v", err) + } + + os.WriteFile(filepath.Join(tmpDir, "test.secret"), []byte("secret"), 0644) + + _, err := w.ReadFile(ctx, "test.secret") + if err == nil { + t.Error("expected error for denied pattern *.secret") + } +} + +func TestHostWorkspace_Exists(t *testing.T) { + tmpDir := t.TempDir() + + w := NewHostWorkspace(tmpDir) + ctx := context.Background() + + if err := w.Initialize(ctx); err != nil { + t.Fatalf("Initialize failed: %v", err) + } + + exists, err := w.Exists(ctx, "nonexistent.txt") + if err != nil { + t.Fatalf("Exists failed: %v", err) + } + if exists { + t.Error("nonexistent file should not exist") + } + + w.WriteFile(ctx, "exists.txt", []byte("content")) + + exists, err = w.Exists(ctx, "exists.txt") + if err != nil { + t.Fatalf("Exists failed: %v", err) + } + if !exists { + t.Error("created file should exist") + } +} + +func TestHostWorkspace_ListFiles(t *testing.T) { + tmpDir := t.TempDir() + + w := NewHostWorkspace(tmpDir) + ctx := context.Background() + + if err := w.Initialize(ctx); err != nil { + t.Fatalf("Initialize failed: %v", err) + } + + w.WriteFile(ctx, "file1.txt", []byte("1")) + w.WriteFile(ctx, "file2.txt", []byte("2")) + w.WriteFile(ctx, "file3.txt", []byte("3")) + + files, err := w.ListFiles(ctx, ".") + if err != nil { + t.Fatalf("ListFiles failed: %v", err) + } + + if len(files) != 3 { + t.Errorf("ListFiles returned %d files, want 3", len(files)) + } +} + +func TestHostWorkspace_DeleteFile(t *testing.T) { + tmpDir := t.TempDir() + + w := NewHostWorkspace(tmpDir) + ctx := context.Background() + + if err := w.Initialize(ctx); err != nil { + t.Fatalf("Initialize failed: %v", err) + } + + w.WriteFile(ctx, "todelete.txt", []byte("content")) + + if err := w.DeleteFile(ctx, "todelete.txt"); err != nil { + t.Fatalf("DeleteFile failed: %v", err) + } + + exists, _ := w.Exists(ctx, "todelete.txt") + if exists { + t.Error("deleted file should not exist") + } +} + +func TestHostWorkspace_IsInWorkspace(t *testing.T) { + tmpDir := t.TempDir() + + w := NewHostWorkspace(tmpDir) + ctx := context.Background() + w.Initialize(ctx) + + if !w.IsInWorkspace("subdir/file.txt") { + t.Error("relative path should be in workspace") + } + + if w.IsInWorkspace("/etc/passwd") { + t.Error("absolute path outside workspace should not be in workspace") + } +} + +func TestHostWorkspace_NotInitialized(t *testing.T) { + tmpDir := t.TempDir() + + w := NewHostWorkspace(tmpDir) + ctx := context.Background() + + _, err := w.ReadFile(ctx, "test.txt") + if err == nil { + t.Error("expected error when workspace not initialized") + } +} diff --git a/test/harness/integration_test.go b/test/harness/integration_test.go new file mode 100644 index 00000000..23fa6078 --- /dev/null +++ b/test/harness/integration_test.go @@ -0,0 +1,135 @@ +//go:build integration + +package harness_test + +import ( + "context" + "os" + "path/filepath" + "testing" + "time" + + "station/internal/config" + "station/pkg/harness" + "station/pkg/harness/workspace" +) + +func TestIntegration_HarnessConfigLoading(t *testing.T) { + cfg, err := config.Load() + if err != nil { + t.Fatalf("Failed to load config: %v", err) + } + + if cfg.Harness.Workspace.Path == "" { + t.Error("Expected harness workspace path to be set") + } + + if cfg.Harness.Workspace.Mode == "" { + t.Error("Expected harness workspace mode to be set") + } + + t.Logf("Harness config loaded: workspace=%s, mode=%s", + cfg.Harness.Workspace.Path, cfg.Harness.Workspace.Mode) +} + +func TestIntegration_WorkspaceCreation(t *testing.T) { + cfg, err := config.Load() + if err != nil { + t.Fatalf("Failed to load config: %v", err) + } + + configDir := config.GetStationConfigDir() + workspacePath := filepath.Join(configDir, cfg.Harness.Workspace.Path) + + ws, err := workspace.NewHostWorkspace(workspacePath, nil) + if err != nil { + t.Fatalf("Failed to create workspace: %v", err) + } + + testDir := filepath.Join("test-integration", time.Now().Format("20060102-150405")) + if err := ws.MkdirAll(testDir); err != nil { + t.Fatalf("Failed to create test directory: %v", err) + } + + testFile := filepath.Join(testDir, "test.txt") + content := []byte("Integration test content") + if err := ws.WriteFile(testFile, content); err != nil { + t.Fatalf("Failed to write test file: %v", err) + } + + readContent, err := ws.ReadFile(testFile) + if err != nil { + t.Fatalf("Failed to read test file: %v", err) + } + + if string(readContent) != string(content) { + t.Errorf("Content mismatch: got %q, want %q", readContent, content) + } + + if err := ws.Remove(testDir); err != nil { + t.Logf("Warning: Failed to cleanup test directory: %v", err) + } + + t.Logf("Workspace integration test passed: %s", workspacePath) +} + +func TestIntegration_ExecutorWithRealLLM(t *testing.T) { + if os.Getenv("INTEGRATION_LLM_TESTS") != "true" { + t.Skip("Skipping LLM test: set INTEGRATION_LLM_TESTS=true to enable") + } + + cfg, err := config.Load() + if err != nil { + t.Fatalf("Failed to load config: %v", err) + } + + if cfg.AIAPIKey == "" { + t.Skip("Skipping: no AI API key configured") + } + + configDir := config.GetStationConfigDir() + workspacePath := filepath.Join(configDir, cfg.Harness.Workspace.Path) + + ws, err := workspace.NewHostWorkspace(workspacePath, nil) + if err != nil { + t.Fatalf("Failed to create workspace: %v", err) + } + + harnessCfg := harness.DefaultHarnessConfig() + agentCfg := harness.DefaultAgentHarnessConfig() + agentCfg.MaxSteps = 5 + agentCfg.Timeout = 2 * time.Minute + + exec := harness.NewExecutor(harnessCfg, agentCfg, ws) + + ctx, cancel := context.WithTimeout(context.Background(), agentCfg.Timeout) + defer cancel() + + result, err := exec.Execute(ctx, "What is 2+2? Respond with just the number.") + if err != nil { + t.Fatalf("Execution failed: %v", err) + } + + t.Logf("Execution completed: steps=%d, status=%s", result.Steps, result.Status) + t.Logf("Response: %s", result.Response) +} + +func TestIntegration_GitWorkflow(t *testing.T) { + if os.Getenv("INTEGRATION_GIT_TESTS") != "true" { + t.Skip("Skipping git test: set INTEGRATION_GIT_TESTS=true to enable") + } + + cfg, err := config.Load() + if err != nil { + t.Fatalf("Failed to load config: %v", err) + } + + configDir := config.GetStationConfigDir() + workspacePath := filepath.Join(configDir, cfg.Harness.Workspace.Path) + + if _, err := os.Stat(filepath.Join(workspacePath, ".git")); os.IsNotExist(err) { + t.Skipf("Skipping: workspace is not a git repository: %s", workspacePath) + } + + t.Logf("Git workflow test would run against: %s", workspacePath) +} diff --git a/ui/src/App.tsx b/ui/src/App.tsx index 378df305..30826474 100644 --- a/ui/src/App.tsx +++ b/ui/src/App.tsx @@ -3921,6 +3921,63 @@ const SettingsPage = () => { + +
+ + updateSandboxConfig({ docker_image: e.target.value })} + placeholder="ubuntu:22.04" + className="w-full bg-white border border-gray-300 text-gray-900 font-mono text-sm p-2 rounded focus:outline-none focus:border-gray-500 focus:ring-1 focus:ring-gray-500" + /> +

Docker image for sandbox containers

+
+ +
+ +
+
+ + updateSandboxConfig({ + registry_auth: { ...configObj.sandbox?.registry_auth, server_address: e.target.value } + })} + placeholder="ghcr.io or 123456.dkr.ecr.us-east-1.amazonaws.com" + className="w-full bg-white border border-gray-300 text-gray-900 font-mono text-xs p-1.5 rounded focus:outline-none focus:border-gray-500" + /> +
+
+
+ + updateSandboxConfig({ + registry_auth: { ...configObj.sandbox?.registry_auth, username: e.target.value } + })} + placeholder="username" + className="w-full bg-white border border-gray-300 text-gray-900 font-mono text-xs p-1.5 rounded focus:outline-none focus:border-gray-500" + /> +
+
+ + updateSandboxConfig({ + registry_auth: { ...configObj.sandbox?.registry_auth, password: e.target.value } + })} + placeholder="••••••••" + className="w-full bg-white border border-gray-300 text-gray-900 font-mono text-xs p-1.5 rounded focus:outline-none focus:border-gray-500" + /> +
+
+

For ECR/GCR, use identity_token instead. Or set STN_SANDBOX_REGISTRY_* env vars.

+
+
)}