| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 1 | package loop |
| 2 | |
| 3 | import ( |
| Josh Bleecher Snyder | a9b3822 | 2025-04-29 18:05:06 -0700 | [diff] [blame] | 4 | "cmp" |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 5 | "context" |
| Josh Bleecher Snyder | dbe0230 | 2025-04-29 16:44:23 -0700 | [diff] [blame] | 6 | _ "embed" |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 7 | "encoding/json" |
| 8 | "fmt" |
| Josh Bleecher Snyder | 3e2111b | 2025-04-30 17:53:28 +0000 | [diff] [blame^] | 9 | "io" |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 10 | "log/slog" |
| 11 | "net/http" |
| 12 | "os" |
| 13 | "os/exec" |
| 14 | "runtime/debug" |
| 15 | "slices" |
| 16 | "strings" |
| 17 | "sync" |
| 18 | "time" |
| 19 | |
| 20 | "sketch.dev/ant" |
| Josh Bleecher Snyder | 3e2111b | 2025-04-30 17:53:28 +0000 | [diff] [blame^] | 21 | "sketch.dev/browser" |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 22 | "sketch.dev/claudetool" |
| Josh Bleecher Snyder | d499fd6 | 2025-04-30 01:31:29 +0000 | [diff] [blame] | 23 | "sketch.dev/claudetool/bashkit" |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 24 | ) |
| 25 | |
| 26 | const ( |
| 27 | userCancelMessage = "user requested agent to stop handling responses" |
| 28 | ) |
| 29 | |
| 30 | type CodingAgent interface { |
| 31 | // Init initializes an agent inside a docker container. |
| 32 | Init(AgentInit) error |
| 33 | |
| 34 | // Ready returns a channel closed after Init successfully called. |
| 35 | Ready() <-chan struct{} |
| 36 | |
| 37 | // URL reports the HTTP URL of this agent. |
| 38 | URL() string |
| 39 | |
| 40 | // UserMessage enqueues a message to the agent and returns immediately. |
| 41 | UserMessage(ctx context.Context, msg string) |
| 42 | |
| 43 | // WaitForMessage blocks until the agent has a response to give. |
| 44 | // Use AgentMessage.EndOfTurn to help determine if you want to |
| 45 | // drain the agent. |
| 46 | WaitForMessage(ctx context.Context) AgentMessage |
| 47 | |
| 48 | // Loop begins the agent loop returns only when ctx is cancelled. |
| 49 | Loop(ctx context.Context) |
| 50 | |
| Sean McCullough | edc88dc | 2025-04-30 02:55:01 +0000 | [diff] [blame] | 51 | CancelTurn(cause error) |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 52 | |
| 53 | CancelToolUse(toolUseID string, cause error) error |
| 54 | |
| 55 | // Returns a subset of the agent's message history. |
| 56 | Messages(start int, end int) []AgentMessage |
| 57 | |
| 58 | // Returns the current number of messages in the history |
| 59 | MessageCount() int |
| 60 | |
| 61 | TotalUsage() ant.CumulativeUsage |
| 62 | OriginalBudget() ant.Budget |
| 63 | |
| 64 | // WaitForMessageCount returns when the agent has at more than clientMessageCount messages or the context is done. |
| 65 | WaitForMessageCount(ctx context.Context, greaterThan int) |
| 66 | |
| 67 | WorkingDir() string |
| 68 | |
| 69 | // Diff returns a unified diff of changes made since the agent was instantiated. |
| 70 | // If commit is non-nil, it shows the diff for just that specific commit. |
| 71 | Diff(commit *string) (string, error) |
| 72 | |
| 73 | // InitialCommit returns the Git commit hash that was saved when the agent was instantiated. |
| 74 | InitialCommit() string |
| 75 | |
| 76 | // Title returns the current title of the conversation. |
| 77 | Title() string |
| 78 | |
| Josh Bleecher Snyder | 47b1936 | 2025-04-30 01:34:14 +0000 | [diff] [blame] | 79 | // BranchName returns the git branch name for the conversation. |
| 80 | BranchName() string |
| 81 | |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 82 | // OS returns the operating system of the client. |
| 83 | OS() string |
| Philip Zeyliger | 99a9a02 | 2025-04-27 15:15:25 +0000 | [diff] [blame] | 84 | |
| Philip Zeyliger | c72fff5 | 2025-04-29 20:17:54 +0000 | [diff] [blame] | 85 | // SessionID returns the unique session identifier. |
| 86 | SessionID() string |
| 87 | |
| Philip Zeyliger | 99a9a02 | 2025-04-27 15:15:25 +0000 | [diff] [blame] | 88 | // OutstandingLLMCallCount returns the number of outstanding LLM calls. |
| 89 | OutstandingLLMCallCount() int |
| 90 | |
| 91 | // OutstandingToolCalls returns the names of outstanding tool calls. |
| 92 | OutstandingToolCalls() []string |
| Philip Zeyliger | 18532b2 | 2025-04-23 21:11:46 +0000 | [diff] [blame] | 93 | OutsideOS() string |
| 94 | OutsideHostname() string |
| 95 | OutsideWorkingDir() string |
| Philip Zeyliger | d140295 | 2025-04-23 03:54:37 +0000 | [diff] [blame] | 96 | GitOrigin() string |
| Josh Bleecher Snyder | 3e2111b | 2025-04-30 17:53:28 +0000 | [diff] [blame^] | 97 | // OpenBrowser is a best-effort attempt to open a browser at url in outside sketch. |
| 98 | OpenBrowser(url string) |
| Philip Zeyliger | 2c4db09 | 2025-04-28 16:57:50 -0700 | [diff] [blame] | 99 | |
| 100 | // RestartConversation resets the conversation history |
| 101 | RestartConversation(ctx context.Context, rev string, initialPrompt string) error |
| 102 | // SuggestReprompt suggests a re-prompt based on the current conversation. |
| 103 | SuggestReprompt(ctx context.Context) (string, error) |
| 104 | // IsInContainer returns true if the agent is running in a container |
| 105 | IsInContainer() bool |
| 106 | // FirstMessageIndex returns the index of the first message in the current conversation |
| 107 | FirstMessageIndex() int |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 108 | } |
| 109 | |
| 110 | type CodingAgentMessageType string |
| 111 | |
| 112 | const ( |
| 113 | UserMessageType CodingAgentMessageType = "user" |
| 114 | AgentMessageType CodingAgentMessageType = "agent" |
| 115 | ErrorMessageType CodingAgentMessageType = "error" |
| 116 | BudgetMessageType CodingAgentMessageType = "budget" // dedicated for "out of budget" errors |
| 117 | ToolUseMessageType CodingAgentMessageType = "tool" |
| 118 | CommitMessageType CodingAgentMessageType = "commit" // for displaying git commits |
| 119 | AutoMessageType CodingAgentMessageType = "auto" // for automated notifications like autoformatting |
| 120 | |
| 121 | cancelToolUseMessage = "Stop responding to my previous message. Wait for me to ask you something else before attempting to use any more tools." |
| 122 | ) |
| 123 | |
| 124 | type AgentMessage struct { |
| 125 | Type CodingAgentMessageType `json:"type"` |
| 126 | // EndOfTurn indicates that the AI is done working and is ready for the next user input. |
| 127 | EndOfTurn bool `json:"end_of_turn"` |
| 128 | |
| 129 | Content string `json:"content"` |
| 130 | ToolName string `json:"tool_name,omitempty"` |
| 131 | ToolInput string `json:"input,omitempty"` |
| 132 | ToolResult string `json:"tool_result,omitempty"` |
| 133 | ToolError bool `json:"tool_error,omitempty"` |
| 134 | ToolCallId string `json:"tool_call_id,omitempty"` |
| 135 | |
| 136 | // ToolCalls is a list of all tool calls requested in this message (name and input pairs) |
| 137 | ToolCalls []ToolCall `json:"tool_calls,omitempty"` |
| 138 | |
| Sean McCullough | d9f1337 | 2025-04-21 15:08:49 -0700 | [diff] [blame] | 139 | // ToolResponses is a list of all responses to tool calls requested in this message (name and input pairs) |
| 140 | ToolResponses []AgentMessage `json:"toolResponses,omitempty"` |
| 141 | |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 142 | // Commits is a list of git commits for a commit message |
| 143 | Commits []*GitCommit `json:"commits,omitempty"` |
| 144 | |
| 145 | Timestamp time.Time `json:"timestamp"` |
| 146 | ConversationID string `json:"conversation_id"` |
| 147 | ParentConversationID *string `json:"parent_conversation_id,omitempty"` |
| 148 | Usage *ant.Usage `json:"usage,omitempty"` |
| 149 | |
| 150 | // Message timing information |
| 151 | StartTime *time.Time `json:"start_time,omitempty"` |
| 152 | EndTime *time.Time `json:"end_time,omitempty"` |
| 153 | Elapsed *time.Duration `json:"elapsed,omitempty"` |
| 154 | |
| 155 | // Turn duration - the time taken for a complete agent turn |
| 156 | TurnDuration *time.Duration `json:"turnDuration,omitempty"` |
| 157 | |
| 158 | Idx int `json:"idx"` |
| 159 | } |
| 160 | |
| Josh Bleecher Snyder | 50a1d62 | 2025-04-29 09:59:03 -0700 | [diff] [blame] | 161 | // SetConvo sets m.ConversationID and m.ParentConversationID based on convo. |
| 162 | func (m *AgentMessage) SetConvo(convo *ant.Convo) { |
| 163 | if convo == nil { |
| 164 | m.ConversationID = "" |
| 165 | m.ParentConversationID = nil |
| 166 | return |
| 167 | } |
| 168 | m.ConversationID = convo.ID |
| 169 | if convo.Parent != nil { |
| 170 | m.ParentConversationID = &convo.Parent.ID |
| 171 | } |
| 172 | } |
| 173 | |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 174 | // GitCommit represents a single git commit for a commit message |
| 175 | type GitCommit struct { |
| 176 | Hash string `json:"hash"` // Full commit hash |
| 177 | Subject string `json:"subject"` // Commit subject line |
| 178 | Body string `json:"body"` // Full commit message body |
| 179 | PushedBranch string `json:"pushed_branch,omitempty"` // If set, this commit was pushed to this branch |
| 180 | } |
| 181 | |
| 182 | // ToolCall represents a single tool call within an agent message |
| 183 | type ToolCall struct { |
| Sean McCullough | d9f1337 | 2025-04-21 15:08:49 -0700 | [diff] [blame] | 184 | Name string `json:"name"` |
| 185 | Input string `json:"input"` |
| 186 | ToolCallId string `json:"tool_call_id"` |
| 187 | ResultMessage *AgentMessage `json:"result_message,omitempty"` |
| 188 | Args string `json:"args,omitempty"` |
| 189 | Result string `json:"result,omitempty"` |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 190 | } |
| 191 | |
| 192 | func (a *AgentMessage) Attr() slog.Attr { |
| 193 | var attrs []any = []any{ |
| 194 | slog.String("type", string(a.Type)), |
| 195 | } |
| 196 | if a.EndOfTurn { |
| 197 | attrs = append(attrs, slog.Bool("end_of_turn", a.EndOfTurn)) |
| 198 | } |
| 199 | if a.Content != "" { |
| 200 | attrs = append(attrs, slog.String("content", a.Content)) |
| 201 | } |
| 202 | if a.ToolName != "" { |
| 203 | attrs = append(attrs, slog.String("tool_name", a.ToolName)) |
| 204 | } |
| 205 | if a.ToolInput != "" { |
| 206 | attrs = append(attrs, slog.String("tool_input", a.ToolInput)) |
| 207 | } |
| 208 | if a.Elapsed != nil { |
| 209 | attrs = append(attrs, slog.Int64("elapsed", a.Elapsed.Nanoseconds())) |
| 210 | } |
| 211 | if a.TurnDuration != nil { |
| 212 | attrs = append(attrs, slog.Int64("turnDuration", a.TurnDuration.Nanoseconds())) |
| 213 | } |
| 214 | if a.ToolResult != "" { |
| 215 | attrs = append(attrs, slog.String("tool_result", a.ToolResult)) |
| 216 | } |
| 217 | if a.ToolError { |
| 218 | attrs = append(attrs, slog.Bool("tool_error", a.ToolError)) |
| 219 | } |
| 220 | if len(a.ToolCalls) > 0 { |
| 221 | toolCallAttrs := make([]any, 0, len(a.ToolCalls)) |
| 222 | for i, tc := range a.ToolCalls { |
| 223 | toolCallAttrs = append(toolCallAttrs, slog.Group( |
| 224 | fmt.Sprintf("tool_call_%d", i), |
| 225 | slog.String("name", tc.Name), |
| 226 | slog.String("input", tc.Input), |
| 227 | )) |
| 228 | } |
| 229 | attrs = append(attrs, slog.Group("tool_calls", toolCallAttrs...)) |
| 230 | } |
| 231 | if a.ConversationID != "" { |
| 232 | attrs = append(attrs, slog.String("convo_id", a.ConversationID)) |
| 233 | } |
| 234 | if a.ParentConversationID != nil { |
| 235 | attrs = append(attrs, slog.String("parent_convo_id", *a.ParentConversationID)) |
| 236 | } |
| 237 | if a.Usage != nil && !a.Usage.IsZero() { |
| 238 | attrs = append(attrs, a.Usage.Attr()) |
| 239 | } |
| 240 | // TODO: timestamp, convo ids, idx? |
| 241 | return slog.Group("agent_message", attrs...) |
| 242 | } |
| 243 | |
| 244 | func errorMessage(err error) AgentMessage { |
| 245 | // It's somewhat unknowable whether error messages are "end of turn" or not, but it seems like the best approach. |
| 246 | if os.Getenv(("DEBUG")) == "1" { |
| 247 | return AgentMessage{Type: ErrorMessageType, Content: err.Error() + " Stacktrace: " + string(debug.Stack()), EndOfTurn: true} |
| 248 | } |
| 249 | |
| 250 | return AgentMessage{Type: ErrorMessageType, Content: err.Error(), EndOfTurn: true} |
| 251 | } |
| 252 | |
| 253 | func budgetMessage(err error) AgentMessage { |
| 254 | return AgentMessage{Type: BudgetMessageType, Content: err.Error(), EndOfTurn: true} |
| 255 | } |
| 256 | |
| 257 | // ConvoInterface defines the interface for conversation interactions |
| 258 | type ConvoInterface interface { |
| 259 | CumulativeUsage() ant.CumulativeUsage |
| 260 | ResetBudget(ant.Budget) |
| 261 | OverBudget() error |
| 262 | SendMessage(message ant.Message) (*ant.MessageResponse, error) |
| 263 | SendUserTextMessage(s string, otherContents ...ant.Content) (*ant.MessageResponse, error) |
| Philip Zeyliger | 2c4db09 | 2025-04-28 16:57:50 -0700 | [diff] [blame] | 264 | GetID() string |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 265 | ToolResultContents(ctx context.Context, resp *ant.MessageResponse) ([]ant.Content, error) |
| 266 | ToolResultCancelContents(resp *ant.MessageResponse) ([]ant.Content, error) |
| 267 | CancelToolUse(toolUseID string, cause error) error |
| Philip Zeyliger | 2c4db09 | 2025-04-28 16:57:50 -0700 | [diff] [blame] | 268 | SubConvoWithHistory() *ant.Convo |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 269 | } |
| 270 | |
| 271 | type Agent struct { |
| Philip Zeyliger | 2c4db09 | 2025-04-28 16:57:50 -0700 | [diff] [blame] | 272 | convo ConvoInterface |
| 273 | config AgentConfig // config for this agent |
| 274 | workingDir string |
| 275 | repoRoot string // workingDir may be a subdir of repoRoot |
| 276 | url string |
| 277 | firstMessageIndex int // index of the first message in the current conversation |
| 278 | lastHEAD string // hash of the last HEAD that was pushed to the host (only when under docker) |
| 279 | initialCommit string // hash of the Git HEAD when the agent was instantiated or Init() |
| 280 | gitRemoteAddr string // HTTP URL of the host git repo (only when under docker) |
| Josh Bleecher Snyder | 3e2111b | 2025-04-30 17:53:28 +0000 | [diff] [blame^] | 281 | outsideHTTP string // base address of the outside webserver (only when under docker) |
| Philip Zeyliger | 2c4db09 | 2025-04-28 16:57:50 -0700 | [diff] [blame] | 282 | ready chan struct{} // closed when the agent is initialized (only when under docker) |
| 283 | startedAt time.Time |
| 284 | originalBudget ant.Budget |
| 285 | title string |
| 286 | branchName string |
| 287 | codereview *claudetool.CodeReviewer |
| Sean McCullough | 96b60dd | 2025-04-30 09:49:10 -0700 | [diff] [blame] | 288 | // State machine to track agent state |
| 289 | stateMachine *StateMachine |
| Philip Zeyliger | 18532b2 | 2025-04-23 21:11:46 +0000 | [diff] [blame] | 290 | // Outside information |
| 291 | outsideHostname string |
| 292 | outsideOS string |
| 293 | outsideWorkingDir string |
| Philip Zeyliger | d140295 | 2025-04-23 03:54:37 +0000 | [diff] [blame] | 294 | // URL of the git remote 'origin' if it exists |
| 295 | gitOrigin string |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 296 | |
| 297 | // Time when the current turn started (reset at the beginning of InnerLoop) |
| 298 | startOfTurn time.Time |
| 299 | |
| 300 | // Inbox - for messages from the user to the agent. |
| 301 | // sent on by UserMessage |
| 302 | // . e.g. when user types into the chat textarea |
| 303 | // read from by GatherMessages |
| 304 | inbox chan string |
| 305 | |
| 306 | // Outbox |
| 307 | // sent on by pushToOutbox |
| 308 | // via OnToolResult and OnResponse callbacks |
| 309 | // read from by WaitForMessage |
| 310 | // called by termui inside its repl loop. |
| 311 | outbox chan AgentMessage |
| 312 | |
| Sean McCullough | edc88dc | 2025-04-30 02:55:01 +0000 | [diff] [blame] | 313 | // protects cancelTurn |
| 314 | cancelTurnMu sync.Mutex |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 315 | // cancels potentially long-running tool_use calls or chains of them |
| Sean McCullough | edc88dc | 2025-04-30 02:55:01 +0000 | [diff] [blame] | 316 | cancelTurn context.CancelCauseFunc |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 317 | |
| 318 | // protects following |
| 319 | mu sync.Mutex |
| 320 | |
| 321 | // Stores all messages for this agent |
| 322 | history []AgentMessage |
| 323 | |
| 324 | listeners []chan struct{} |
| 325 | |
| 326 | // Track git commits we've already seen (by hash) |
| 327 | seenCommits map[string]bool |
| Philip Zeyliger | 99a9a02 | 2025-04-27 15:15:25 +0000 | [diff] [blame] | 328 | |
| 329 | // Track outstanding LLM call IDs |
| 330 | outstandingLLMCalls map[string]struct{} |
| 331 | |
| 332 | // Track outstanding tool calls by ID with their names |
| 333 | outstandingToolCalls map[string]string |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 334 | } |
| 335 | |
| 336 | func (a *Agent) URL() string { return a.url } |
| 337 | |
| 338 | // Title returns the current title of the conversation. |
| 339 | // If no title has been set, returns an empty string. |
| 340 | func (a *Agent) Title() string { |
| 341 | a.mu.Lock() |
| 342 | defer a.mu.Unlock() |
| 343 | return a.title |
| 344 | } |
| 345 | |
| Josh Bleecher Snyder | 47b1936 | 2025-04-30 01:34:14 +0000 | [diff] [blame] | 346 | // BranchName returns the git branch name for the conversation. |
| 347 | func (a *Agent) BranchName() string { |
| 348 | a.mu.Lock() |
| 349 | defer a.mu.Unlock() |
| 350 | return a.branchName |
| 351 | } |
| 352 | |
| Philip Zeyliger | 99a9a02 | 2025-04-27 15:15:25 +0000 | [diff] [blame] | 353 | // OutstandingLLMCallCount returns the number of outstanding LLM calls. |
| 354 | func (a *Agent) OutstandingLLMCallCount() int { |
| 355 | a.mu.Lock() |
| 356 | defer a.mu.Unlock() |
| 357 | return len(a.outstandingLLMCalls) |
| 358 | } |
| 359 | |
| 360 | // OutstandingToolCalls returns the names of outstanding tool calls. |
| 361 | func (a *Agent) OutstandingToolCalls() []string { |
| 362 | a.mu.Lock() |
| 363 | defer a.mu.Unlock() |
| 364 | |
| 365 | tools := make([]string, 0, len(a.outstandingToolCalls)) |
| 366 | for _, toolName := range a.outstandingToolCalls { |
| 367 | tools = append(tools, toolName) |
| 368 | } |
| 369 | return tools |
| 370 | } |
| 371 | |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 372 | // OS returns the operating system of the client. |
| 373 | func (a *Agent) OS() string { |
| 374 | return a.config.ClientGOOS |
| 375 | } |
| 376 | |
| Philip Zeyliger | c72fff5 | 2025-04-29 20:17:54 +0000 | [diff] [blame] | 377 | func (a *Agent) SessionID() string { |
| 378 | return a.config.SessionID |
| 379 | } |
| 380 | |
| Philip Zeyliger | 18532b2 | 2025-04-23 21:11:46 +0000 | [diff] [blame] | 381 | // OutsideOS returns the operating system of the outside system. |
| 382 | func (a *Agent) OutsideOS() string { |
| 383 | return a.outsideOS |
| Philip Zeyliger | d140295 | 2025-04-23 03:54:37 +0000 | [diff] [blame] | 384 | } |
| 385 | |
| Philip Zeyliger | 18532b2 | 2025-04-23 21:11:46 +0000 | [diff] [blame] | 386 | // OutsideHostname returns the hostname of the outside system. |
| 387 | func (a *Agent) OutsideHostname() string { |
| 388 | return a.outsideHostname |
| Philip Zeyliger | d140295 | 2025-04-23 03:54:37 +0000 | [diff] [blame] | 389 | } |
| 390 | |
| Philip Zeyliger | 18532b2 | 2025-04-23 21:11:46 +0000 | [diff] [blame] | 391 | // OutsideWorkingDir returns the working directory on the outside system. |
| 392 | func (a *Agent) OutsideWorkingDir() string { |
| 393 | return a.outsideWorkingDir |
| Philip Zeyliger | d140295 | 2025-04-23 03:54:37 +0000 | [diff] [blame] | 394 | } |
| 395 | |
| 396 | // GitOrigin returns the URL of the git remote 'origin' if it exists. |
| 397 | func (a *Agent) GitOrigin() string { |
| 398 | return a.gitOrigin |
| 399 | } |
| 400 | |
| Josh Bleecher Snyder | 3e2111b | 2025-04-30 17:53:28 +0000 | [diff] [blame^] | 401 | func (a *Agent) OpenBrowser(url string) { |
| 402 | if !a.IsInContainer() { |
| 403 | browser.Open(url) |
| 404 | return |
| 405 | } |
| 406 | // We're in Docker, need to send a request to the Git server |
| 407 | // to signal that the outer process should open the browser. |
| 408 | httpc := &http.Client{Timeout: 5 * time.Second} |
| 409 | resp, err := httpc.Post(a.outsideHTTP+"/browser", "text/plain", strings.NewReader(url)) |
| 410 | if err != nil { |
| 411 | slog.Debug("browser launch request connection failed", "err", err, "url", url) |
| 412 | return |
| 413 | } |
| 414 | defer resp.Body.Close() |
| 415 | if resp.StatusCode == http.StatusOK { |
| 416 | return |
| 417 | } |
| 418 | body, _ := io.ReadAll(resp.Body) |
| 419 | slog.Debug("browser launch request execution failed", "status", resp.Status, "body", string(body)) |
| 420 | } |
| 421 | |
| Sean McCullough | 96b60dd | 2025-04-30 09:49:10 -0700 | [diff] [blame] | 422 | // CurrentState returns the current state of the agent's state machine. |
| 423 | func (a *Agent) CurrentState() State { |
| 424 | return a.stateMachine.CurrentState() |
| 425 | } |
| 426 | |
| Philip Zeyliger | 2c4db09 | 2025-04-28 16:57:50 -0700 | [diff] [blame] | 427 | func (a *Agent) IsInContainer() bool { |
| 428 | return a.config.InDocker |
| 429 | } |
| 430 | |
| 431 | func (a *Agent) FirstMessageIndex() int { |
| 432 | a.mu.Lock() |
| 433 | defer a.mu.Unlock() |
| 434 | return a.firstMessageIndex |
| 435 | } |
| 436 | |
| Josh Bleecher Snyder | a9b3822 | 2025-04-29 18:05:06 -0700 | [diff] [blame] | 437 | // SetTitleBranch sets the title and branch name of the conversation. |
| 438 | func (a *Agent) SetTitleBranch(title, branchName string) { |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 439 | a.mu.Lock() |
| 440 | defer a.mu.Unlock() |
| 441 | a.title = title |
| Josh Bleecher Snyder | a9b3822 | 2025-04-29 18:05:06 -0700 | [diff] [blame] | 442 | a.branchName = branchName |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 443 | // Notify all listeners that the state has changed |
| 444 | for _, ch := range a.listeners { |
| 445 | close(ch) |
| 446 | } |
| 447 | a.listeners = a.listeners[:0] |
| 448 | } |
| 449 | |
| Philip Zeyliger | 99a9a02 | 2025-04-27 15:15:25 +0000 | [diff] [blame] | 450 | // OnToolCall implements ant.Listener and tracks the start of a tool call. |
| 451 | func (a *Agent) OnToolCall(ctx context.Context, convo *ant.Convo, id string, toolName string, toolInput json.RawMessage, content ant.Content) { |
| 452 | // Track the tool call |
| 453 | a.mu.Lock() |
| 454 | a.outstandingToolCalls[id] = toolName |
| 455 | a.mu.Unlock() |
| 456 | } |
| 457 | |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 458 | // OnToolResult implements ant.Listener. |
| Philip Zeyliger | 99a9a02 | 2025-04-27 15:15:25 +0000 | [diff] [blame] | 459 | func (a *Agent) OnToolResult(ctx context.Context, convo *ant.Convo, toolID string, toolName string, toolInput json.RawMessage, content ant.Content, result *string, err error) { |
| 460 | // Remove the tool call from outstanding calls |
| 461 | a.mu.Lock() |
| 462 | delete(a.outstandingToolCalls, toolID) |
| 463 | a.mu.Unlock() |
| 464 | |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 465 | m := AgentMessage{ |
| 466 | Type: ToolUseMessageType, |
| 467 | Content: content.Text, |
| 468 | ToolResult: content.ToolResult, |
| 469 | ToolError: content.ToolError, |
| 470 | ToolName: toolName, |
| 471 | ToolInput: string(toolInput), |
| 472 | ToolCallId: content.ToolUseID, |
| 473 | StartTime: content.StartTime, |
| 474 | EndTime: content.EndTime, |
| 475 | } |
| 476 | |
| 477 | // Calculate the elapsed time if both start and end times are set |
| 478 | if content.StartTime != nil && content.EndTime != nil { |
| 479 | elapsed := content.EndTime.Sub(*content.StartTime) |
| 480 | m.Elapsed = &elapsed |
| 481 | } |
| 482 | |
| Josh Bleecher Snyder | 50a1d62 | 2025-04-29 09:59:03 -0700 | [diff] [blame] | 483 | m.SetConvo(convo) |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 484 | a.pushToOutbox(ctx, m) |
| 485 | } |
| 486 | |
| 487 | // OnRequest implements ant.Listener. |
| Philip Zeyliger | 99a9a02 | 2025-04-27 15:15:25 +0000 | [diff] [blame] | 488 | func (a *Agent) OnRequest(ctx context.Context, convo *ant.Convo, id string, msg *ant.Message) { |
| 489 | a.mu.Lock() |
| 490 | defer a.mu.Unlock() |
| 491 | a.outstandingLLMCalls[id] = struct{}{} |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 492 | // We already get tool results from the above. We send user messages to the outbox in the agent loop. |
| 493 | } |
| 494 | |
| 495 | // OnResponse implements ant.Listener. Responses contain messages from the LLM |
| 496 | // that need to be displayed (as well as tool calls that we send along when |
| 497 | // they're done). (It would be reasonable to also mention tool calls when they're |
| 498 | // started, but we don't do that yet.) |
| Philip Zeyliger | 99a9a02 | 2025-04-27 15:15:25 +0000 | [diff] [blame] | 499 | func (a *Agent) OnResponse(ctx context.Context, convo *ant.Convo, id string, resp *ant.MessageResponse) { |
| 500 | // Remove the LLM call from outstanding calls |
| 501 | a.mu.Lock() |
| 502 | delete(a.outstandingLLMCalls, id) |
| 503 | a.mu.Unlock() |
| 504 | |
| Josh Bleecher Snyder | 50a1d62 | 2025-04-29 09:59:03 -0700 | [diff] [blame] | 505 | if resp == nil { |
| 506 | // LLM API call failed |
| 507 | m := AgentMessage{ |
| 508 | Type: ErrorMessageType, |
| 509 | Content: "API call failed, type 'continue' to try again", |
| 510 | } |
| 511 | m.SetConvo(convo) |
| 512 | a.pushToOutbox(ctx, m) |
| 513 | return |
| 514 | } |
| 515 | |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 516 | endOfTurn := false |
| 517 | if resp.StopReason != ant.StopReasonToolUse { |
| 518 | endOfTurn = true |
| 519 | } |
| 520 | m := AgentMessage{ |
| 521 | Type: AgentMessageType, |
| 522 | Content: collectTextContent(resp), |
| 523 | EndOfTurn: endOfTurn, |
| 524 | Usage: &resp.Usage, |
| 525 | StartTime: resp.StartTime, |
| 526 | EndTime: resp.EndTime, |
| 527 | } |
| 528 | |
| 529 | // Extract any tool calls from the response |
| 530 | if resp.StopReason == ant.StopReasonToolUse { |
| 531 | var toolCalls []ToolCall |
| 532 | for _, part := range resp.Content { |
| 533 | if part.Type == "tool_use" { |
| 534 | toolCalls = append(toolCalls, ToolCall{ |
| 535 | Name: part.ToolName, |
| 536 | Input: string(part.ToolInput), |
| 537 | ToolCallId: part.ID, |
| 538 | }) |
| 539 | } |
| 540 | } |
| 541 | m.ToolCalls = toolCalls |
| 542 | } |
| 543 | |
| 544 | // Calculate the elapsed time if both start and end times are set |
| 545 | if resp.StartTime != nil && resp.EndTime != nil { |
| 546 | elapsed := resp.EndTime.Sub(*resp.StartTime) |
| 547 | m.Elapsed = &elapsed |
| 548 | } |
| 549 | |
| Josh Bleecher Snyder | 50a1d62 | 2025-04-29 09:59:03 -0700 | [diff] [blame] | 550 | m.SetConvo(convo) |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 551 | a.pushToOutbox(ctx, m) |
| 552 | } |
| 553 | |
| 554 | // WorkingDir implements CodingAgent. |
| 555 | func (a *Agent) WorkingDir() string { |
| 556 | return a.workingDir |
| 557 | } |
| 558 | |
| 559 | // MessageCount implements CodingAgent. |
| 560 | func (a *Agent) MessageCount() int { |
| 561 | a.mu.Lock() |
| 562 | defer a.mu.Unlock() |
| 563 | return len(a.history) |
| 564 | } |
| 565 | |
| 566 | // Messages implements CodingAgent. |
| 567 | func (a *Agent) Messages(start int, end int) []AgentMessage { |
| 568 | a.mu.Lock() |
| 569 | defer a.mu.Unlock() |
| 570 | return slices.Clone(a.history[start:end]) |
| 571 | } |
| 572 | |
| 573 | func (a *Agent) OriginalBudget() ant.Budget { |
| 574 | return a.originalBudget |
| 575 | } |
| 576 | |
| 577 | // AgentConfig contains configuration for creating a new Agent. |
| 578 | type AgentConfig struct { |
| 579 | Context context.Context |
| 580 | AntURL string |
| 581 | APIKey string |
| 582 | HTTPC *http.Client |
| 583 | Budget ant.Budget |
| 584 | GitUsername string |
| 585 | GitEmail string |
| 586 | SessionID string |
| 587 | ClientGOOS string |
| 588 | ClientGOARCH string |
| Philip Zeyliger | 2c4db09 | 2025-04-28 16:57:50 -0700 | [diff] [blame] | 589 | InDocker bool |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 590 | UseAnthropicEdit bool |
| Philip Zeyliger | 18532b2 | 2025-04-23 21:11:46 +0000 | [diff] [blame] | 591 | // Outside information |
| 592 | OutsideHostname string |
| 593 | OutsideOS string |
| 594 | OutsideWorkingDir string |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 595 | } |
| 596 | |
| 597 | // NewAgent creates a new Agent. |
| 598 | // It is not usable until Init() is called. |
| 599 | func NewAgent(config AgentConfig) *Agent { |
| 600 | agent := &Agent{ |
| Philip Zeyliger | 99a9a02 | 2025-04-27 15:15:25 +0000 | [diff] [blame] | 601 | config: config, |
| 602 | ready: make(chan struct{}), |
| 603 | inbox: make(chan string, 100), |
| 604 | outbox: make(chan AgentMessage, 100), |
| 605 | startedAt: time.Now(), |
| 606 | originalBudget: config.Budget, |
| 607 | seenCommits: make(map[string]bool), |
| 608 | outsideHostname: config.OutsideHostname, |
| 609 | outsideOS: config.OutsideOS, |
| 610 | outsideWorkingDir: config.OutsideWorkingDir, |
| 611 | outstandingLLMCalls: make(map[string]struct{}), |
| 612 | outstandingToolCalls: make(map[string]string), |
| Sean McCullough | 96b60dd | 2025-04-30 09:49:10 -0700 | [diff] [blame] | 613 | stateMachine: NewStateMachine(), |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 614 | } |
| 615 | return agent |
| 616 | } |
| 617 | |
| 618 | type AgentInit struct { |
| 619 | WorkingDir string |
| 620 | NoGit bool // only for testing |
| 621 | |
| 622 | InDocker bool |
| 623 | Commit string |
| Josh Bleecher Snyder | 3e2111b | 2025-04-30 17:53:28 +0000 | [diff] [blame^] | 624 | OutsideHTTP string |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 625 | GitRemoteAddr string |
| 626 | HostAddr string |
| 627 | } |
| 628 | |
| 629 | func (a *Agent) Init(ini AgentInit) error { |
| Josh Bleecher Snyder | 9c07e1d | 2025-04-28 19:25:37 -0700 | [diff] [blame] | 630 | if a.convo != nil { |
| 631 | return fmt.Errorf("Agent.Init: already initialized") |
| 632 | } |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 633 | ctx := a.config.Context |
| 634 | if ini.InDocker { |
| 635 | cmd := exec.CommandContext(ctx, "git", "stash") |
| 636 | cmd.Dir = ini.WorkingDir |
| 637 | if out, err := cmd.CombinedOutput(); err != nil { |
| 638 | return fmt.Errorf("git stash: %s: %v", out, err) |
| 639 | } |
| Philip Zeyliger | d0ac1ea | 2025-04-21 20:04:19 -0700 | [diff] [blame] | 640 | cmd = exec.CommandContext(ctx, "git", "remote", "add", "sketch-host", ini.GitRemoteAddr) |
| 641 | cmd.Dir = ini.WorkingDir |
| 642 | if out, err := cmd.CombinedOutput(); err != nil { |
| 643 | return fmt.Errorf("git remote add: %s: %v", out, err) |
| 644 | } |
| 645 | cmd = exec.CommandContext(ctx, "git", "fetch", "sketch-host") |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 646 | cmd.Dir = ini.WorkingDir |
| 647 | if out, err := cmd.CombinedOutput(); err != nil { |
| 648 | return fmt.Errorf("git fetch: %s: %w", out, err) |
| 649 | } |
| 650 | cmd = exec.CommandContext(ctx, "git", "checkout", "-f", ini.Commit) |
| 651 | cmd.Dir = ini.WorkingDir |
| 652 | if out, err := cmd.CombinedOutput(); err != nil { |
| 653 | return fmt.Errorf("git checkout %s: %s: %w", ini.Commit, out, err) |
| 654 | } |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 655 | a.lastHEAD = ini.Commit |
| 656 | a.gitRemoteAddr = ini.GitRemoteAddr |
| Josh Bleecher Snyder | 3e2111b | 2025-04-30 17:53:28 +0000 | [diff] [blame^] | 657 | a.outsideHTTP = ini.OutsideHTTP |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 658 | a.initialCommit = ini.Commit |
| 659 | if ini.HostAddr != "" { |
| 660 | a.url = "http://" + ini.HostAddr |
| 661 | } |
| 662 | } |
| 663 | a.workingDir = ini.WorkingDir |
| 664 | |
| 665 | if !ini.NoGit { |
| 666 | repoRoot, err := repoRoot(ctx, a.workingDir) |
| 667 | if err != nil { |
| 668 | return fmt.Errorf("repoRoot: %w", err) |
| 669 | } |
| 670 | a.repoRoot = repoRoot |
| 671 | |
| 672 | commitHash, err := resolveRef(ctx, a.repoRoot, "HEAD") |
| 673 | if err != nil { |
| 674 | return fmt.Errorf("resolveRef: %w", err) |
| 675 | } |
| 676 | a.initialCommit = commitHash |
| 677 | |
| 678 | codereview, err := claudetool.NewCodeReviewer(ctx, a.repoRoot, a.initialCommit) |
| 679 | if err != nil { |
| 680 | return fmt.Errorf("Agent.Init: claudetool.NewCodeReviewer: %w", err) |
| 681 | } |
| 682 | a.codereview = codereview |
| Philip Zeyliger | d140295 | 2025-04-23 03:54:37 +0000 | [diff] [blame] | 683 | |
| 684 | a.gitOrigin = getGitOrigin(ctx, ini.WorkingDir) |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 685 | } |
| 686 | a.lastHEAD = a.initialCommit |
| 687 | a.convo = a.initConvo() |
| 688 | close(a.ready) |
| 689 | return nil |
| 690 | } |
| 691 | |
| Josh Bleecher Snyder | dbe0230 | 2025-04-29 16:44:23 -0700 | [diff] [blame] | 692 | //go:embed agent_system_prompt.txt |
| 693 | var agentSystemPrompt string |
| 694 | |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 695 | // initConvo initializes the conversation. |
| 696 | // It must not be called until all agent fields are initialized, |
| 697 | // particularly workingDir and git. |
| 698 | func (a *Agent) initConvo() *ant.Convo { |
| 699 | ctx := a.config.Context |
| 700 | convo := ant.NewConvo(ctx, a.config.APIKey) |
| 701 | if a.config.HTTPC != nil { |
| 702 | convo.HTTPC = a.config.HTTPC |
| 703 | } |
| 704 | if a.config.AntURL != "" { |
| 705 | convo.URL = a.config.AntURL |
| 706 | } |
| 707 | convo.PromptCaching = true |
| 708 | convo.Budget = a.config.Budget |
| 709 | |
| 710 | var editPrompt string |
| 711 | if a.config.UseAnthropicEdit { |
| 712 | editPrompt = "Then use the str_replace_editor tool to make those edits. For short complete file replacements, you may use the bash tool with cat and heredoc stdin." |
| 713 | } else { |
| 714 | editPrompt = "Then use the patch tool to make those edits. Combine all edits to any given file into a single patch tool call." |
| 715 | } |
| 716 | |
| Josh Bleecher Snyder | dbe0230 | 2025-04-29 16:44:23 -0700 | [diff] [blame] | 717 | convo.SystemPrompt = fmt.Sprintf(agentSystemPrompt, editPrompt, a.config.ClientGOOS, a.config.ClientGOARCH, a.workingDir, a.repoRoot, a.initialCommit) |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 718 | |
| Josh Bleecher Snyder | d499fd6 | 2025-04-30 01:31:29 +0000 | [diff] [blame] | 719 | // Define a permission callback for the bash tool to check if the branch name is set before allowing git commits |
| 720 | bashPermissionCheck := func(command string) error { |
| 721 | // Check if branch name is set |
| 722 | a.mu.Lock() |
| 723 | branchSet := a.branchName != "" |
| 724 | a.mu.Unlock() |
| 725 | |
| 726 | // If branch is set, all commands are allowed |
| 727 | if branchSet { |
| 728 | return nil |
| 729 | } |
| 730 | |
| 731 | // If branch is not set, check if this is a git commit command |
| 732 | willCommit, err := bashkit.WillRunGitCommit(command) |
| 733 | if err != nil { |
| 734 | // If there's an error checking, we should allow the command to proceed |
| 735 | return nil |
| 736 | } |
| 737 | |
| 738 | // If it's a git commit and branch is not set, return an error |
| 739 | if willCommit { |
| 740 | return fmt.Errorf("you must use the title tool before making git commits") |
| 741 | } |
| 742 | |
| 743 | return nil |
| 744 | } |
| 745 | |
| 746 | // Create a custom bash tool with the permission check |
| 747 | bashTool := claudetool.NewBashTool(bashPermissionCheck) |
| 748 | |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 749 | // Register all tools with the conversation |
| 750 | // When adding, removing, or modifying tools here, double-check that the termui tool display |
| 751 | // template in termui/termui.go has pretty-printing support for all tools. |
| 752 | convo.Tools = []*ant.Tool{ |
| Josh Bleecher Snyder | d499fd6 | 2025-04-30 01:31:29 +0000 | [diff] [blame] | 753 | bashTool, claudetool.Keyword, |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 754 | claudetool.Think, a.titleTool(), makeDoneTool(a.codereview, a.config.GitUsername, a.config.GitEmail), |
| 755 | a.codereview.Tool(), |
| 756 | } |
| 757 | if a.config.UseAnthropicEdit { |
| 758 | convo.Tools = append(convo.Tools, claudetool.AnthropicEditTool) |
| 759 | } else { |
| 760 | convo.Tools = append(convo.Tools, claudetool.Patch) |
| 761 | } |
| 762 | convo.Listener = a |
| 763 | return convo |
| 764 | } |
| 765 | |
| Josh Bleecher Snyder | fff269b | 2025-04-30 01:49:39 +0000 | [diff] [blame] | 766 | // branchExists reports whether branchName exists, either locally or in well-known remotes. |
| 767 | func branchExists(dir, branchName string) bool { |
| 768 | refs := []string{ |
| 769 | "refs/heads/", |
| 770 | "refs/remotes/origin/", |
| 771 | "refs/remotes/sketch-host/", |
| 772 | } |
| 773 | for _, ref := range refs { |
| 774 | cmd := exec.Command("git", "show-ref", "--verify", "--quiet", ref+branchName) |
| 775 | cmd.Dir = dir |
| 776 | if cmd.Run() == nil { // exit code 0 means branch exists |
| 777 | return true |
| 778 | } |
| 779 | } |
| 780 | return false |
| 781 | } |
| 782 | |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 783 | func (a *Agent) titleTool() *ant.Tool { |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 784 | title := &ant.Tool{ |
| 785 | Name: "title", |
| Josh Bleecher Snyder | a9b3822 | 2025-04-29 18:05:06 -0700 | [diff] [blame] | 786 | Description: `Sets the conversation title and creates a git branch for tracking work. MANDATORY: You must use this tool before making any git commits.`, |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 787 | InputSchema: json.RawMessage(`{ |
| 788 | "type": "object", |
| 789 | "properties": { |
| 790 | "title": { |
| 791 | "type": "string", |
| Josh Bleecher Snyder | 250348e | 2025-04-30 10:31:28 -0700 | [diff] [blame] | 792 | "description": "A concise title summarizing what this conversation is about, imperative tense preferred" |
| Josh Bleecher Snyder | a9b3822 | 2025-04-29 18:05:06 -0700 | [diff] [blame] | 793 | }, |
| 794 | "branch_name": { |
| 795 | "type": "string", |
| 796 | "description": "A 2-3 word alphanumeric hyphenated slug for the git branch name" |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 797 | } |
| 798 | }, |
| Josh Bleecher Snyder | a9b3822 | 2025-04-29 18:05:06 -0700 | [diff] [blame] | 799 | "required": ["title", "branch_name"] |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 800 | }`), |
| 801 | Run: func(ctx context.Context, input json.RawMessage) (string, error) { |
| 802 | var params struct { |
| Josh Bleecher Snyder | a9b3822 | 2025-04-29 18:05:06 -0700 | [diff] [blame] | 803 | Title string `json:"title"` |
| 804 | BranchName string `json:"branch_name"` |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 805 | } |
| 806 | if err := json.Unmarshal(input, ¶ms); err != nil { |
| 807 | return "", err |
| 808 | } |
| Josh Bleecher Snyder | a9b3822 | 2025-04-29 18:05:06 -0700 | [diff] [blame] | 809 | // It's unfortunate to not allow title changes, |
| 810 | // but it avoids having multiple branches. |
| 811 | t := a.Title() |
| 812 | if t != "" { |
| 813 | return "", fmt.Errorf("title already set to: %s", t) |
| 814 | } |
| 815 | |
| 816 | if params.BranchName == "" { |
| 817 | return "", fmt.Errorf("branch_name parameter cannot be empty") |
| 818 | } |
| 819 | if params.Title == "" { |
| 820 | return "", fmt.Errorf("title parameter cannot be empty") |
| 821 | } |
| Josh Bleecher Snyder | 42f7a7c | 2025-04-30 10:29:21 -0700 | [diff] [blame] | 822 | if params.BranchName != cleanBranchName(params.BranchName) { |
| 823 | return "", fmt.Errorf("branch_name parameter must be alphanumeric hyphenated slug") |
| 824 | } |
| 825 | branchName := "sketch/" + params.BranchName |
| Josh Bleecher Snyder | fff269b | 2025-04-30 01:49:39 +0000 | [diff] [blame] | 826 | if branchExists(a.workingDir, branchName) { |
| 827 | return "", fmt.Errorf("branch %q already exists; please choose a different branch name", branchName) |
| 828 | } |
| 829 | |
| Josh Bleecher Snyder | a9b3822 | 2025-04-29 18:05:06 -0700 | [diff] [blame] | 830 | a.SetTitleBranch(params.Title, branchName) |
| 831 | |
| 832 | response := fmt.Sprintf("Title set to %q, branch name set to %q", params.Title, branchName) |
| 833 | return response, nil |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 834 | }, |
| 835 | } |
| 836 | return title |
| 837 | } |
| 838 | |
| 839 | func (a *Agent) Ready() <-chan struct{} { |
| 840 | return a.ready |
| 841 | } |
| 842 | |
| 843 | func (a *Agent) UserMessage(ctx context.Context, msg string) { |
| 844 | a.pushToOutbox(ctx, AgentMessage{Type: UserMessageType, Content: msg}) |
| 845 | a.inbox <- msg |
| 846 | } |
| 847 | |
| 848 | func (a *Agent) WaitForMessage(ctx context.Context) AgentMessage { |
| 849 | // TODO: Should this drain any outbox messages in case there are multiple? |
| 850 | select { |
| 851 | case msg := <-a.outbox: |
| 852 | return msg |
| 853 | case <-ctx.Done(): |
| 854 | return errorMessage(ctx.Err()) |
| 855 | } |
| 856 | } |
| 857 | |
| 858 | func (a *Agent) CancelToolUse(toolUseID string, cause error) error { |
| 859 | return a.convo.CancelToolUse(toolUseID, cause) |
| 860 | } |
| 861 | |
| Sean McCullough | edc88dc | 2025-04-30 02:55:01 +0000 | [diff] [blame] | 862 | func (a *Agent) CancelTurn(cause error) { |
| 863 | a.cancelTurnMu.Lock() |
| 864 | defer a.cancelTurnMu.Unlock() |
| 865 | if a.cancelTurn != nil { |
| Sean McCullough | 96b60dd | 2025-04-30 09:49:10 -0700 | [diff] [blame] | 866 | // Force state transition to cancelled state |
| 867 | ctx := a.config.Context |
| 868 | a.stateMachine.ForceTransition(ctx, StateCancelled, "User cancelled turn: "+cause.Error()) |
| Sean McCullough | edc88dc | 2025-04-30 02:55:01 +0000 | [diff] [blame] | 869 | a.cancelTurn(cause) |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 870 | } |
| 871 | } |
| 872 | |
| 873 | func (a *Agent) Loop(ctxOuter context.Context) { |
| 874 | for { |
| 875 | select { |
| 876 | case <-ctxOuter.Done(): |
| 877 | return |
| 878 | default: |
| 879 | ctxInner, cancel := context.WithCancelCause(ctxOuter) |
| Sean McCullough | edc88dc | 2025-04-30 02:55:01 +0000 | [diff] [blame] | 880 | a.cancelTurnMu.Lock() |
| 881 | // Set .cancelTurn so the user can cancel whatever is happening |
| Sean McCullough | 885a16a | 2025-04-30 02:49:25 +0000 | [diff] [blame] | 882 | // inside the conversation loop without canceling this outer Loop execution. |
| Sean McCullough | edc88dc | 2025-04-30 02:55:01 +0000 | [diff] [blame] | 883 | // This cancelTurn func is intended be called from other goroutines, |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 884 | // hence the mutex. |
| Sean McCullough | edc88dc | 2025-04-30 02:55:01 +0000 | [diff] [blame] | 885 | a.cancelTurn = cancel |
| 886 | a.cancelTurnMu.Unlock() |
| Sean McCullough | 9f4b808 | 2025-04-30 17:34:07 +0000 | [diff] [blame] | 887 | err := a.processTurn(ctxInner) // Renamed from InnerLoop to better reflect its purpose |
| 888 | if err != nil { |
| 889 | slog.ErrorContext(ctxOuter, "Error in processing turn", "error", err) |
| 890 | } |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 891 | cancel(nil) |
| 892 | } |
| 893 | } |
| 894 | } |
| 895 | |
| 896 | func (a *Agent) pushToOutbox(ctx context.Context, m AgentMessage) { |
| 897 | if m.Timestamp.IsZero() { |
| 898 | m.Timestamp = time.Now() |
| 899 | } |
| 900 | |
| 901 | // If this is an end-of-turn message, calculate the turn duration and add it to the message |
| 902 | if m.EndOfTurn && m.Type == AgentMessageType { |
| 903 | turnDuration := time.Since(a.startOfTurn) |
| 904 | m.TurnDuration = &turnDuration |
| 905 | slog.InfoContext(ctx, "Turn completed", "turnDuration", turnDuration) |
| 906 | } |
| 907 | |
| 908 | slog.InfoContext(ctx, "agent message", m.Attr()) |
| 909 | |
| 910 | a.mu.Lock() |
| 911 | defer a.mu.Unlock() |
| 912 | m.Idx = len(a.history) |
| 913 | a.history = append(a.history, m) |
| 914 | a.outbox <- m |
| 915 | |
| 916 | // Notify all listeners: |
| 917 | for _, ch := range a.listeners { |
| 918 | close(ch) |
| 919 | } |
| 920 | a.listeners = a.listeners[:0] |
| 921 | } |
| 922 | |
| 923 | func (a *Agent) GatherMessages(ctx context.Context, block bool) ([]ant.Content, error) { |
| 924 | var m []ant.Content |
| 925 | if block { |
| 926 | select { |
| 927 | case <-ctx.Done(): |
| 928 | return m, ctx.Err() |
| 929 | case msg := <-a.inbox: |
| 930 | m = append(m, ant.Content{Type: "text", Text: msg}) |
| 931 | } |
| 932 | } |
| 933 | for { |
| 934 | select { |
| 935 | case msg := <-a.inbox: |
| 936 | m = append(m, ant.Content{Type: "text", Text: msg}) |
| 937 | default: |
| 938 | return m, nil |
| 939 | } |
| 940 | } |
| 941 | } |
| 942 | |
| Sean McCullough | 885a16a | 2025-04-30 02:49:25 +0000 | [diff] [blame] | 943 | // processTurn handles a single conversation turn with the user |
| Sean McCullough | 9f4b808 | 2025-04-30 17:34:07 +0000 | [diff] [blame] | 944 | func (a *Agent) processTurn(ctx context.Context) error { |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 945 | // Reset the start of turn time |
| 946 | a.startOfTurn = time.Now() |
| 947 | |
| Sean McCullough | 96b60dd | 2025-04-30 09:49:10 -0700 | [diff] [blame] | 948 | // Transition to waiting for user input state |
| 949 | a.stateMachine.Transition(ctx, StateWaitingForUserInput, "Starting turn") |
| 950 | |
| Sean McCullough | 885a16a | 2025-04-30 02:49:25 +0000 | [diff] [blame] | 951 | // Process initial user message |
| 952 | initialResp, err := a.processUserMessage(ctx) |
| 953 | if err != nil { |
| Sean McCullough | 96b60dd | 2025-04-30 09:49:10 -0700 | [diff] [blame] | 954 | a.stateMachine.Transition(ctx, StateError, "Error processing user message: "+err.Error()) |
| Sean McCullough | 9f4b808 | 2025-04-30 17:34:07 +0000 | [diff] [blame] | 955 | return err |
| 956 | } |
| 957 | |
| 958 | // Handle edge case where both initialResp and err are nil |
| 959 | if initialResp == nil { |
| 960 | err := fmt.Errorf("unexpected nil response from processUserMessage with no error") |
| Sean McCullough | 96b60dd | 2025-04-30 09:49:10 -0700 | [diff] [blame] | 961 | a.stateMachine.Transition(ctx, StateError, "Error processing user message: "+err.Error()) |
| 962 | |
| Sean McCullough | 9f4b808 | 2025-04-30 17:34:07 +0000 | [diff] [blame] | 963 | a.pushToOutbox(ctx, errorMessage(err)) |
| 964 | return err |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 965 | } |
| Sean McCullough | 885a16a | 2025-04-30 02:49:25 +0000 | [diff] [blame] | 966 | |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 967 | // We do this as we go, but let's also do it at the end of the turn |
| 968 | defer func() { |
| 969 | if _, err := a.handleGitCommits(ctx); err != nil { |
| 970 | // Just log the error, don't stop execution |
| 971 | slog.WarnContext(ctx, "Failed to check for new git commits", "error", err) |
| 972 | } |
| 973 | }() |
| 974 | |
| Sean McCullough | 885a16a | 2025-04-30 02:49:25 +0000 | [diff] [blame] | 975 | // Main response loop - continue as long as the model is using tools |
| 976 | resp := initialResp |
| 977 | for { |
| 978 | // Check if we are over budget |
| 979 | if err := a.overBudget(ctx); err != nil { |
| Sean McCullough | 96b60dd | 2025-04-30 09:49:10 -0700 | [diff] [blame] | 980 | a.stateMachine.Transition(ctx, StateBudgetExceeded, "Budget exceeded: "+err.Error()) |
| Sean McCullough | 9f4b808 | 2025-04-30 17:34:07 +0000 | [diff] [blame] | 981 | return err |
| Sean McCullough | 885a16a | 2025-04-30 02:49:25 +0000 | [diff] [blame] | 982 | } |
| 983 | |
| 984 | // If the model is not requesting to use a tool, we're done |
| 985 | if resp.StopReason != ant.StopReasonToolUse { |
| Sean McCullough | 96b60dd | 2025-04-30 09:49:10 -0700 | [diff] [blame] | 986 | a.stateMachine.Transition(ctx, StateEndOfTurn, "LLM completed response, ending turn") |
| Sean McCullough | 885a16a | 2025-04-30 02:49:25 +0000 | [diff] [blame] | 987 | break |
| 988 | } |
| 989 | |
| Sean McCullough | 96b60dd | 2025-04-30 09:49:10 -0700 | [diff] [blame] | 990 | // Transition to tool use requested state |
| 991 | a.stateMachine.Transition(ctx, StateToolUseRequested, "LLM requested tool use") |
| 992 | |
| Sean McCullough | 885a16a | 2025-04-30 02:49:25 +0000 | [diff] [blame] | 993 | // Handle tool execution |
| 994 | continueConversation, toolResp := a.handleToolExecution(ctx, resp) |
| 995 | if !continueConversation { |
| Sean McCullough | 9f4b808 | 2025-04-30 17:34:07 +0000 | [diff] [blame] | 996 | return nil |
| Sean McCullough | 885a16a | 2025-04-30 02:49:25 +0000 | [diff] [blame] | 997 | } |
| 998 | |
| 999 | // Set the response for the next iteration |
| 1000 | resp = toolResp |
| 1001 | } |
| Sean McCullough | 9f4b808 | 2025-04-30 17:34:07 +0000 | [diff] [blame] | 1002 | |
| 1003 | return nil |
| Sean McCullough | 885a16a | 2025-04-30 02:49:25 +0000 | [diff] [blame] | 1004 | } |
| 1005 | |
| 1006 | // processUserMessage waits for user messages and sends them to the model |
| 1007 | func (a *Agent) processUserMessage(ctx context.Context) (*ant.MessageResponse, error) { |
| 1008 | // Wait for at least one message from the user |
| 1009 | msgs, err := a.GatherMessages(ctx, true) |
| 1010 | if err != nil { // e.g. the context was canceled while blocking in GatherMessages |
| Sean McCullough | 96b60dd | 2025-04-30 09:49:10 -0700 | [diff] [blame] | 1011 | a.stateMachine.Transition(ctx, StateError, "Error gathering messages: "+err.Error()) |
| Sean McCullough | 885a16a | 2025-04-30 02:49:25 +0000 | [diff] [blame] | 1012 | return nil, err |
| 1013 | } |
| 1014 | |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 1015 | userMessage := ant.Message{ |
| 1016 | Role: "user", |
| 1017 | Content: msgs, |
| 1018 | } |
| Sean McCullough | 885a16a | 2025-04-30 02:49:25 +0000 | [diff] [blame] | 1019 | |
| Sean McCullough | 96b60dd | 2025-04-30 09:49:10 -0700 | [diff] [blame] | 1020 | // Transition to sending to LLM state |
| 1021 | a.stateMachine.Transition(ctx, StateSendingToLLM, "Sending user message to LLM") |
| 1022 | |
| Sean McCullough | 885a16a | 2025-04-30 02:49:25 +0000 | [diff] [blame] | 1023 | // Send message to the model |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 1024 | resp, err := a.convo.SendMessage(userMessage) |
| 1025 | if err != nil { |
| Sean McCullough | 96b60dd | 2025-04-30 09:49:10 -0700 | [diff] [blame] | 1026 | a.stateMachine.Transition(ctx, StateError, "Error sending to LLM: "+err.Error()) |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 1027 | a.pushToOutbox(ctx, errorMessage(err)) |
| Sean McCullough | 885a16a | 2025-04-30 02:49:25 +0000 | [diff] [blame] | 1028 | return nil, err |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 1029 | } |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 1030 | |
| Sean McCullough | 96b60dd | 2025-04-30 09:49:10 -0700 | [diff] [blame] | 1031 | // Transition to processing LLM response state |
| 1032 | a.stateMachine.Transition(ctx, StateProcessingLLMResponse, "Processing LLM response") |
| 1033 | |
| Sean McCullough | 885a16a | 2025-04-30 02:49:25 +0000 | [diff] [blame] | 1034 | return resp, nil |
| 1035 | } |
| 1036 | |
| 1037 | // handleToolExecution processes a tool use request from the model |
| 1038 | func (a *Agent) handleToolExecution(ctx context.Context, resp *ant.MessageResponse) (bool, *ant.MessageResponse) { |
| 1039 | var results []ant.Content |
| 1040 | cancelled := false |
| 1041 | |
| Sean McCullough | 96b60dd | 2025-04-30 09:49:10 -0700 | [diff] [blame] | 1042 | // Transition to checking for cancellation state |
| 1043 | a.stateMachine.Transition(ctx, StateCheckingForCancellation, "Checking if user requested cancellation") |
| 1044 | |
| Sean McCullough | 885a16a | 2025-04-30 02:49:25 +0000 | [diff] [blame] | 1045 | // Check if the operation was cancelled by the user |
| 1046 | select { |
| 1047 | case <-ctx.Done(): |
| 1048 | // Don't actually run any of the tools, but rather build a response |
| 1049 | // for each tool_use message letting the LLM know that user canceled it. |
| 1050 | var err error |
| 1051 | results, err = a.convo.ToolResultCancelContents(resp) |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 1052 | if err != nil { |
| Sean McCullough | 96b60dd | 2025-04-30 09:49:10 -0700 | [diff] [blame] | 1053 | a.stateMachine.Transition(ctx, StateError, "Error creating cancellation response: "+err.Error()) |
| Sean McCullough | 885a16a | 2025-04-30 02:49:25 +0000 | [diff] [blame] | 1054 | a.pushToOutbox(ctx, errorMessage(err)) |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 1055 | } |
| Sean McCullough | 885a16a | 2025-04-30 02:49:25 +0000 | [diff] [blame] | 1056 | cancelled = true |
| Sean McCullough | 96b60dd | 2025-04-30 09:49:10 -0700 | [diff] [blame] | 1057 | a.stateMachine.Transition(ctx, StateCancelled, "Operation cancelled by user") |
| Sean McCullough | 885a16a | 2025-04-30 02:49:25 +0000 | [diff] [blame] | 1058 | default: |
| Sean McCullough | 96b60dd | 2025-04-30 09:49:10 -0700 | [diff] [blame] | 1059 | // Transition to running tool state |
| 1060 | a.stateMachine.Transition(ctx, StateRunningTool, "Executing requested tool") |
| 1061 | |
| Sean McCullough | 885a16a | 2025-04-30 02:49:25 +0000 | [diff] [blame] | 1062 | // Add working directory to context for tool execution |
| 1063 | ctx = claudetool.WithWorkingDir(ctx, a.workingDir) |
| 1064 | |
| 1065 | // Execute the tools |
| 1066 | var err error |
| 1067 | results, err = a.convo.ToolResultContents(ctx, resp) |
| 1068 | if ctx.Err() != nil { // e.g. the user canceled the operation |
| 1069 | cancelled = true |
| Sean McCullough | 96b60dd | 2025-04-30 09:49:10 -0700 | [diff] [blame] | 1070 | a.stateMachine.Transition(ctx, StateCancelled, "Operation cancelled during tool execution") |
| Sean McCullough | 885a16a | 2025-04-30 02:49:25 +0000 | [diff] [blame] | 1071 | } else if err != nil { |
| Sean McCullough | 96b60dd | 2025-04-30 09:49:10 -0700 | [diff] [blame] | 1072 | a.stateMachine.Transition(ctx, StateError, "Error executing tool: "+err.Error()) |
| Sean McCullough | 885a16a | 2025-04-30 02:49:25 +0000 | [diff] [blame] | 1073 | a.pushToOutbox(ctx, errorMessage(err)) |
| 1074 | } |
| 1075 | } |
| 1076 | |
| 1077 | // Process git commits that may have occurred during tool execution |
| Sean McCullough | 96b60dd | 2025-04-30 09:49:10 -0700 | [diff] [blame] | 1078 | a.stateMachine.Transition(ctx, StateCheckingGitCommits, "Checking for git commits") |
| Sean McCullough | 885a16a | 2025-04-30 02:49:25 +0000 | [diff] [blame] | 1079 | autoqualityMessages := a.processGitChanges(ctx) |
| 1080 | |
| 1081 | // Check budget again after tool execution |
| Sean McCullough | 96b60dd | 2025-04-30 09:49:10 -0700 | [diff] [blame] | 1082 | a.stateMachine.Transition(ctx, StateCheckingBudget, "Checking budget after tool execution") |
| Sean McCullough | 885a16a | 2025-04-30 02:49:25 +0000 | [diff] [blame] | 1083 | if err := a.overBudget(ctx); err != nil { |
| Sean McCullough | 96b60dd | 2025-04-30 09:49:10 -0700 | [diff] [blame] | 1084 | a.stateMachine.Transition(ctx, StateBudgetExceeded, "Budget exceeded after tool execution: "+err.Error()) |
| Sean McCullough | 885a16a | 2025-04-30 02:49:25 +0000 | [diff] [blame] | 1085 | return false, nil |
| 1086 | } |
| 1087 | |
| 1088 | // Continue the conversation with tool results and any user messages |
| 1089 | return a.continueTurnWithToolResults(ctx, results, autoqualityMessages, cancelled) |
| 1090 | } |
| 1091 | |
| 1092 | // processGitChanges checks for new git commits and runs autoformatters if needed |
| 1093 | func (a *Agent) processGitChanges(ctx context.Context) []string { |
| 1094 | // Check for git commits after tool execution |
| 1095 | newCommits, err := a.handleGitCommits(ctx) |
| 1096 | if err != nil { |
| 1097 | // Just log the error, don't stop execution |
| 1098 | slog.WarnContext(ctx, "Failed to check for new git commits", "error", err) |
| 1099 | return nil |
| 1100 | } |
| 1101 | |
| 1102 | // Run autoformatters if there was exactly one new commit |
| 1103 | var autoqualityMessages []string |
| 1104 | if len(newCommits) == 1 { |
| Sean McCullough | 96b60dd | 2025-04-30 09:49:10 -0700 | [diff] [blame] | 1105 | a.stateMachine.Transition(ctx, StateRunningAutoformatters, "Running autoformatters on new commit") |
| Sean McCullough | 885a16a | 2025-04-30 02:49:25 +0000 | [diff] [blame] | 1106 | formatted := a.codereview.Autoformat(ctx) |
| 1107 | if len(formatted) > 0 { |
| 1108 | msg := fmt.Sprintf(` |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 1109 | I ran autoformatters and they updated these files: |
| 1110 | |
| 1111 | %s |
| 1112 | |
| 1113 | Please amend your latest git commit with these changes and then continue with what you were doing.`, |
| Sean McCullough | 885a16a | 2025-04-30 02:49:25 +0000 | [diff] [blame] | 1114 | strings.Join(formatted, "\n"), |
| 1115 | )[1:] |
| 1116 | a.pushToOutbox(ctx, AgentMessage{ |
| 1117 | Type: AutoMessageType, |
| 1118 | Content: msg, |
| 1119 | Timestamp: time.Now(), |
| 1120 | }) |
| 1121 | autoqualityMessages = append(autoqualityMessages, msg) |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 1122 | } |
| 1123 | } |
| Sean McCullough | 885a16a | 2025-04-30 02:49:25 +0000 | [diff] [blame] | 1124 | |
| 1125 | return autoqualityMessages |
| 1126 | } |
| 1127 | |
| 1128 | // continueTurnWithToolResults continues the conversation with tool results |
| 1129 | func (a *Agent) continueTurnWithToolResults(ctx context.Context, results []ant.Content, autoqualityMessages []string, cancelled bool) (bool, *ant.MessageResponse) { |
| 1130 | // Get any messages the user sent while tools were executing |
| Sean McCullough | 96b60dd | 2025-04-30 09:49:10 -0700 | [diff] [blame] | 1131 | a.stateMachine.Transition(ctx, StateGatheringAdditionalMessages, "Gathering additional user messages") |
| Sean McCullough | 885a16a | 2025-04-30 02:49:25 +0000 | [diff] [blame] | 1132 | msgs, err := a.GatherMessages(ctx, false) |
| 1133 | if err != nil { |
| Sean McCullough | 96b60dd | 2025-04-30 09:49:10 -0700 | [diff] [blame] | 1134 | a.stateMachine.Transition(ctx, StateError, "Error gathering additional messages: "+err.Error()) |
| Sean McCullough | 885a16a | 2025-04-30 02:49:25 +0000 | [diff] [blame] | 1135 | return false, nil |
| 1136 | } |
| 1137 | |
| 1138 | // Inject any auto-generated messages from quality checks |
| 1139 | for _, msg := range autoqualityMessages { |
| 1140 | msgs = append(msgs, ant.Content{Type: "text", Text: msg}) |
| 1141 | } |
| 1142 | |
| 1143 | // Handle cancellation by appending a message about it |
| 1144 | if cancelled { |
| 1145 | msgs = append(msgs, ant.Content{Type: "text", Text: cancelToolUseMessage}) |
| 1146 | // EndOfTurn is false here so that the client of this agent keeps processing |
| 1147 | // messages from WaitForMessage() and gets the response from the LLM |
| 1148 | a.pushToOutbox(ctx, AgentMessage{Type: ErrorMessageType, Content: userCancelMessage, EndOfTurn: false}) |
| 1149 | } else if err := a.convo.OverBudget(); err != nil { |
| 1150 | // Handle budget issues by appending a message about it |
| 1151 | budgetMsg := "We've exceeded our budget. Please ask the user to confirm before continuing by ending the turn." |
| 1152 | msgs = append(msgs, ant.Content{Type: "text", Text: budgetMsg}) |
| 1153 | a.pushToOutbox(ctx, budgetMessage(fmt.Errorf("warning: %w (ask to keep trying, if you'd like)", err))) |
| 1154 | } |
| 1155 | |
| 1156 | // Combine tool results with user messages |
| 1157 | results = append(results, msgs...) |
| 1158 | |
| 1159 | // Send the combined message to continue the conversation |
| Sean McCullough | 96b60dd | 2025-04-30 09:49:10 -0700 | [diff] [blame] | 1160 | a.stateMachine.Transition(ctx, StateSendingToolResults, "Sending tool results back to LLM") |
| Sean McCullough | 885a16a | 2025-04-30 02:49:25 +0000 | [diff] [blame] | 1161 | resp, err := a.convo.SendMessage(ant.Message{ |
| 1162 | Role: "user", |
| 1163 | Content: results, |
| 1164 | }) |
| 1165 | if err != nil { |
| Sean McCullough | 96b60dd | 2025-04-30 09:49:10 -0700 | [diff] [blame] | 1166 | a.stateMachine.Transition(ctx, StateError, "Error sending tool results: "+err.Error()) |
| Sean McCullough | 885a16a | 2025-04-30 02:49:25 +0000 | [diff] [blame] | 1167 | a.pushToOutbox(ctx, errorMessage(fmt.Errorf("error: failed to continue conversation: %s", err.Error()))) |
| 1168 | return true, nil // Return true to continue the conversation, but with no response |
| 1169 | } |
| 1170 | |
| Sean McCullough | 96b60dd | 2025-04-30 09:49:10 -0700 | [diff] [blame] | 1171 | // Transition back to processing LLM response |
| 1172 | a.stateMachine.Transition(ctx, StateProcessingLLMResponse, "Processing LLM response to tool results") |
| 1173 | |
| Sean McCullough | 885a16a | 2025-04-30 02:49:25 +0000 | [diff] [blame] | 1174 | if cancelled { |
| 1175 | return false, nil |
| 1176 | } |
| 1177 | |
| 1178 | return true, resp |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 1179 | } |
| 1180 | |
| 1181 | func (a *Agent) overBudget(ctx context.Context) error { |
| 1182 | if err := a.convo.OverBudget(); err != nil { |
| Sean McCullough | 96b60dd | 2025-04-30 09:49:10 -0700 | [diff] [blame] | 1183 | a.stateMachine.Transition(ctx, StateBudgetExceeded, "Budget exceeded: "+err.Error()) |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 1184 | m := budgetMessage(err) |
| 1185 | m.Content = m.Content + "\n\nBudget reset." |
| 1186 | a.pushToOutbox(ctx, budgetMessage(err)) |
| 1187 | a.convo.ResetBudget(a.originalBudget) |
| 1188 | return err |
| 1189 | } |
| 1190 | return nil |
| 1191 | } |
| 1192 | |
| 1193 | func collectTextContent(msg *ant.MessageResponse) string { |
| 1194 | // Collect all text content |
| 1195 | var allText strings.Builder |
| 1196 | for _, content := range msg.Content { |
| 1197 | if content.Type == "text" && content.Text != "" { |
| 1198 | if allText.Len() > 0 { |
| 1199 | allText.WriteString("\n\n") |
| 1200 | } |
| 1201 | allText.WriteString(content.Text) |
| 1202 | } |
| 1203 | } |
| 1204 | return allText.String() |
| 1205 | } |
| 1206 | |
| 1207 | func (a *Agent) TotalUsage() ant.CumulativeUsage { |
| 1208 | a.mu.Lock() |
| 1209 | defer a.mu.Unlock() |
| 1210 | return a.convo.CumulativeUsage() |
| 1211 | } |
| 1212 | |
| 1213 | // WaitForMessageCount returns when the agent has at more than clientMessageCount messages or the context is done. |
| 1214 | func (a *Agent) WaitForMessageCount(ctx context.Context, greaterThan int) { |
| 1215 | for a.MessageCount() <= greaterThan { |
| 1216 | a.mu.Lock() |
| 1217 | ch := make(chan struct{}) |
| 1218 | // Deletion happens when we notify. |
| 1219 | a.listeners = append(a.listeners, ch) |
| 1220 | a.mu.Unlock() |
| 1221 | |
| 1222 | select { |
| 1223 | case <-ctx.Done(): |
| 1224 | return |
| 1225 | case <-ch: |
| 1226 | continue |
| 1227 | } |
| 1228 | } |
| 1229 | } |
| 1230 | |
| 1231 | // Diff returns a unified diff of changes made since the agent was instantiated. |
| 1232 | func (a *Agent) Diff(commit *string) (string, error) { |
| 1233 | if a.initialCommit == "" { |
| 1234 | return "", fmt.Errorf("no initial commit reference available") |
| 1235 | } |
| 1236 | |
| 1237 | // Find the repository root |
| 1238 | ctx := context.Background() |
| 1239 | |
| 1240 | // If a specific commit hash is provided, show just that commit's changes |
| 1241 | if commit != nil && *commit != "" { |
| 1242 | // Validate that the commit looks like a valid git SHA |
| 1243 | if !isValidGitSHA(*commit) { |
| 1244 | return "", fmt.Errorf("invalid git commit SHA format: %s", *commit) |
| 1245 | } |
| 1246 | |
| 1247 | // Get the diff for just this commit |
| 1248 | cmd := exec.CommandContext(ctx, "git", "show", "--unified=10", *commit) |
| 1249 | cmd.Dir = a.repoRoot |
| 1250 | output, err := cmd.CombinedOutput() |
| 1251 | if err != nil { |
| 1252 | return "", fmt.Errorf("failed to get diff for commit %s: %w - %s", *commit, err, string(output)) |
| 1253 | } |
| 1254 | return string(output), nil |
| 1255 | } |
| 1256 | |
| 1257 | // Otherwise, get the diff between the initial commit and the current state using exec.Command |
| 1258 | cmd := exec.CommandContext(ctx, "git", "diff", "--unified=10", a.initialCommit) |
| 1259 | cmd.Dir = a.repoRoot |
| 1260 | output, err := cmd.CombinedOutput() |
| 1261 | if err != nil { |
| 1262 | return "", fmt.Errorf("failed to get diff: %w - %s", err, string(output)) |
| 1263 | } |
| 1264 | |
| 1265 | return string(output), nil |
| 1266 | } |
| 1267 | |
| 1268 | // InitialCommit returns the Git commit hash that was saved when the agent was instantiated. |
| 1269 | func (a *Agent) InitialCommit() string { |
| 1270 | return a.initialCommit |
| 1271 | } |
| 1272 | |
| 1273 | // handleGitCommits() highlights new commits to the user. When running |
| 1274 | // under docker, new HEADs are pushed to a branch according to the title. |
| 1275 | func (a *Agent) handleGitCommits(ctx context.Context) ([]*GitCommit, error) { |
| 1276 | if a.repoRoot == "" { |
| 1277 | return nil, nil |
| 1278 | } |
| 1279 | |
| 1280 | head, err := resolveRef(ctx, a.repoRoot, "HEAD") |
| 1281 | if err != nil { |
| 1282 | return nil, err |
| 1283 | } |
| 1284 | if head == a.lastHEAD { |
| 1285 | return nil, nil // nothing to do |
| 1286 | } |
| 1287 | defer func() { |
| 1288 | a.lastHEAD = head |
| 1289 | }() |
| 1290 | |
| 1291 | // Get new commits. Because it's possible that the agent does rebases, fixups, and |
| 1292 | // so forth, we use, as our fixed point, the "initialCommit", and we limit ourselves |
| 1293 | // to the last 100 commits. |
| 1294 | var commits []*GitCommit |
| 1295 | |
| 1296 | // Get commits since the initial commit |
| 1297 | // Format: <hash>\0<subject>\0<body>\0 |
| 1298 | // This uses NULL bytes as separators to avoid issues with newlines in commit messages |
| 1299 | // Limit to 100 commits to avoid overwhelming the user |
| 1300 | cmd := exec.CommandContext(ctx, "git", "log", "-n", "100", "--pretty=format:%H%x00%s%x00%b%x00", "^"+a.initialCommit, head) |
| 1301 | cmd.Dir = a.repoRoot |
| 1302 | output, err := cmd.Output() |
| 1303 | if err != nil { |
| 1304 | return nil, fmt.Errorf("failed to get git log: %w", err) |
| 1305 | } |
| 1306 | |
| 1307 | // Parse git log output and filter out already seen commits |
| 1308 | parsedCommits := parseGitLog(string(output)) |
| 1309 | |
| 1310 | var headCommit *GitCommit |
| 1311 | |
| 1312 | // Filter out commits we've already seen |
| 1313 | for _, commit := range parsedCommits { |
| 1314 | if commit.Hash == head { |
| 1315 | headCommit = &commit |
| 1316 | } |
| 1317 | |
| 1318 | // Skip if we've seen this commit before. If our head has changed, always include that. |
| 1319 | if a.seenCommits[commit.Hash] && commit.Hash != head { |
| 1320 | continue |
| 1321 | } |
| 1322 | |
| 1323 | // Mark this commit as seen |
| 1324 | a.seenCommits[commit.Hash] = true |
| 1325 | |
| 1326 | // Add to our list of new commits |
| 1327 | commits = append(commits, &commit) |
| 1328 | } |
| 1329 | |
| 1330 | if a.gitRemoteAddr != "" { |
| 1331 | if headCommit == nil { |
| 1332 | // I think this can only happen if we have a bug or if there's a race. |
| 1333 | headCommit = &GitCommit{} |
| 1334 | headCommit.Hash = head |
| 1335 | headCommit.Subject = "unknown" |
| 1336 | commits = append(commits, headCommit) |
| 1337 | } |
| 1338 | |
| Josh Bleecher Snyder | a9b3822 | 2025-04-29 18:05:06 -0700 | [diff] [blame] | 1339 | branch := cmp.Or(a.branchName, "sketch/"+a.config.SessionID) |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 1340 | |
| 1341 | // TODO: I don't love the force push here. We could see if the push is a fast-forward, and, |
| 1342 | // if it's not, we could make a backup with a unique name (perhaps append a timestamp) and |
| 1343 | // then use push with lease to replace. |
| 1344 | cmd = exec.Command("git", "push", "--force", a.gitRemoteAddr, "HEAD:refs/heads/"+branch) |
| 1345 | cmd.Dir = a.workingDir |
| 1346 | if out, err := cmd.CombinedOutput(); err != nil { |
| 1347 | a.pushToOutbox(ctx, errorMessage(fmt.Errorf("git push to host: %s: %v", out, err))) |
| 1348 | } else { |
| 1349 | headCommit.PushedBranch = branch |
| 1350 | } |
| 1351 | } |
| 1352 | |
| 1353 | // If we found new commits, create a message |
| 1354 | if len(commits) > 0 { |
| 1355 | msg := AgentMessage{ |
| 1356 | Type: CommitMessageType, |
| 1357 | Timestamp: time.Now(), |
| 1358 | Commits: commits, |
| 1359 | } |
| 1360 | a.pushToOutbox(ctx, msg) |
| 1361 | } |
| 1362 | return commits, nil |
| 1363 | } |
| 1364 | |
| Josh Bleecher Snyder | a9b3822 | 2025-04-29 18:05:06 -0700 | [diff] [blame] | 1365 | func cleanBranchName(s string) string { |
| Josh Bleecher Snyder | 1ae976b | 2025-04-30 00:06:43 +0000 | [diff] [blame] | 1366 | return strings.Map(func(r rune) rune { |
| 1367 | // lowercase |
| 1368 | if r >= 'A' && r <= 'Z' { |
| 1369 | return r + 'a' - 'A' |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 1370 | } |
| Josh Bleecher Snyder | 1ae976b | 2025-04-30 00:06:43 +0000 | [diff] [blame] | 1371 | // replace spaces with dashes |
| 1372 | if r == ' ' { |
| 1373 | return '-' |
| 1374 | } |
| 1375 | // allow alphanumerics and dashes |
| 1376 | if (r >= 'a' && r <= 'z') || r == '-' || (r >= '0' && r <= '9') { |
| 1377 | return r |
| 1378 | } |
| 1379 | return -1 |
| 1380 | }, s) |
| Earl Lee | 2e463fb | 2025-04-17 11:22:22 -0700 | [diff] [blame] | 1381 | } |
| 1382 | |
| 1383 | // parseGitLog parses the output of git log with format '%H%x00%s%x00%b%x00' |
| 1384 | // and returns an array of GitCommit structs. |
| 1385 | func parseGitLog(output string) []GitCommit { |
| 1386 | var commits []GitCommit |
| 1387 | |
| 1388 | // No output means no commits |
| 1389 | if len(output) == 0 { |
| 1390 | return commits |
| 1391 | } |
| 1392 | |
| 1393 | // Split by NULL byte |
| 1394 | parts := strings.Split(output, "\x00") |
| 1395 | |
| 1396 | // Process in triplets (hash, subject, body) |
| 1397 | for i := 0; i < len(parts); i++ { |
| 1398 | // Skip empty parts |
| 1399 | if parts[i] == "" { |
| 1400 | continue |
| 1401 | } |
| 1402 | |
| 1403 | // This should be a hash |
| 1404 | hash := strings.TrimSpace(parts[i]) |
| 1405 | |
| 1406 | // Make sure we have at least a subject part available |
| 1407 | if i+1 >= len(parts) { |
| 1408 | break // No more parts available |
| 1409 | } |
| 1410 | |
| 1411 | // Get the subject |
| 1412 | subject := strings.TrimSpace(parts[i+1]) |
| 1413 | |
| 1414 | // Get the body if available |
| 1415 | body := "" |
| 1416 | if i+2 < len(parts) { |
| 1417 | body = strings.TrimSpace(parts[i+2]) |
| 1418 | } |
| 1419 | |
| 1420 | // Skip to the next triplet |
| 1421 | i += 2 |
| 1422 | |
| 1423 | commits = append(commits, GitCommit{ |
| 1424 | Hash: hash, |
| 1425 | Subject: subject, |
| 1426 | Body: body, |
| 1427 | }) |
| 1428 | } |
| 1429 | |
| 1430 | return commits |
| 1431 | } |
| 1432 | |
| 1433 | func repoRoot(ctx context.Context, dir string) (string, error) { |
| 1434 | cmd := exec.CommandContext(ctx, "git", "rev-parse", "--show-toplevel") |
| 1435 | stderr := new(strings.Builder) |
| 1436 | cmd.Stderr = stderr |
| 1437 | cmd.Dir = dir |
| 1438 | out, err := cmd.Output() |
| 1439 | if err != nil { |
| 1440 | return "", fmt.Errorf("git rev-parse failed: %w\n%s", err, stderr) |
| 1441 | } |
| 1442 | return strings.TrimSpace(string(out)), nil |
| 1443 | } |
| 1444 | |
| 1445 | func resolveRef(ctx context.Context, dir, refName string) (string, error) { |
| 1446 | cmd := exec.CommandContext(ctx, "git", "rev-parse", refName) |
| 1447 | stderr := new(strings.Builder) |
| 1448 | cmd.Stderr = stderr |
| 1449 | cmd.Dir = dir |
| 1450 | out, err := cmd.Output() |
| 1451 | if err != nil { |
| 1452 | return "", fmt.Errorf("git rev-parse failed: %w\n%s", err, stderr) |
| 1453 | } |
| 1454 | // TODO: validate that out is valid hex |
| 1455 | return strings.TrimSpace(string(out)), nil |
| 1456 | } |
| 1457 | |
| 1458 | // isValidGitSHA validates if a string looks like a valid git SHA hash. |
| 1459 | // Git SHAs are hexadecimal strings of at least 4 characters but typically 7, 8, or 40 characters. |
| 1460 | func isValidGitSHA(sha string) bool { |
| 1461 | // Git SHA must be a hexadecimal string with at least 4 characters |
| 1462 | if len(sha) < 4 || len(sha) > 40 { |
| 1463 | return false |
| 1464 | } |
| 1465 | |
| 1466 | // Check if the string only contains hexadecimal characters |
| 1467 | for _, char := range sha { |
| 1468 | if !(char >= '0' && char <= '9') && !(char >= 'a' && char <= 'f') && !(char >= 'A' && char <= 'F') { |
| 1469 | return false |
| 1470 | } |
| 1471 | } |
| 1472 | |
| 1473 | return true |
| 1474 | } |
| Philip Zeyliger | d140295 | 2025-04-23 03:54:37 +0000 | [diff] [blame] | 1475 | |
| 1476 | // getGitOrigin returns the URL of the git remote 'origin' if it exists |
| 1477 | func getGitOrigin(ctx context.Context, dir string) string { |
| 1478 | cmd := exec.CommandContext(ctx, "git", "config", "--get", "remote.origin.url") |
| 1479 | cmd.Dir = dir |
| 1480 | stderr := new(strings.Builder) |
| 1481 | cmd.Stderr = stderr |
| 1482 | out, err := cmd.Output() |
| 1483 | if err != nil { |
| 1484 | return "" |
| 1485 | } |
| 1486 | return strings.TrimSpace(string(out)) |
| 1487 | } |
| Philip Zeyliger | 2c4db09 | 2025-04-28 16:57:50 -0700 | [diff] [blame] | 1488 | |
| 1489 | func (a *Agent) initGitRevision(ctx context.Context, workingDir, revision string) error { |
| 1490 | cmd := exec.CommandContext(ctx, "git", "stash") |
| 1491 | cmd.Dir = workingDir |
| 1492 | if out, err := cmd.CombinedOutput(); err != nil { |
| 1493 | return fmt.Errorf("git stash: %s: %v", out, err) |
| 1494 | } |
| 1495 | cmd = exec.CommandContext(ctx, "git", "fetch", "sketch-host") |
| 1496 | cmd.Dir = workingDir |
| 1497 | if out, err := cmd.CombinedOutput(); err != nil { |
| 1498 | return fmt.Errorf("git fetch: %s: %w", out, err) |
| 1499 | } |
| 1500 | cmd = exec.CommandContext(ctx, "git", "checkout", "-f", revision) |
| 1501 | cmd.Dir = workingDir |
| 1502 | if out, err := cmd.CombinedOutput(); err != nil { |
| 1503 | return fmt.Errorf("git checkout %s: %s: %w", revision, out, err) |
| 1504 | } |
| 1505 | a.lastHEAD = revision |
| 1506 | a.initialCommit = revision |
| 1507 | return nil |
| 1508 | } |
| 1509 | |
| 1510 | func (a *Agent) RestartConversation(ctx context.Context, rev string, initialPrompt string) error { |
| 1511 | a.mu.Lock() |
| 1512 | a.title = "" |
| 1513 | a.firstMessageIndex = len(a.history) |
| 1514 | a.convo = a.initConvo() |
| 1515 | gitReset := func() error { |
| 1516 | if a.config.InDocker && rev != "" { |
| 1517 | err := a.initGitRevision(ctx, a.workingDir, rev) |
| 1518 | if err != nil { |
| 1519 | return err |
| 1520 | } |
| 1521 | } else if !a.config.InDocker && rev != "" { |
| 1522 | return fmt.Errorf("Not resetting git repo when working outside of a container.") |
| 1523 | } |
| 1524 | return nil |
| 1525 | } |
| 1526 | err := gitReset() |
| 1527 | a.mu.Unlock() |
| 1528 | if err != nil { |
| 1529 | a.pushToOutbox(a.config.Context, errorMessage(err)) |
| 1530 | } |
| 1531 | |
| 1532 | a.pushToOutbox(a.config.Context, AgentMessage{ |
| 1533 | Type: AgentMessageType, Content: "Conversation restarted.", |
| 1534 | }) |
| 1535 | if initialPrompt != "" { |
| 1536 | a.UserMessage(ctx, initialPrompt) |
| 1537 | } |
| 1538 | return nil |
| 1539 | } |
| 1540 | |
| 1541 | func (a *Agent) SuggestReprompt(ctx context.Context) (string, error) { |
| 1542 | msg := `The user has requested a suggestion for a re-prompt. |
| 1543 | |
| 1544 | Given the current conversation thus far, suggest a re-prompt that would |
| 1545 | capture the instructions and feedback so far, as well as any |
| 1546 | research or other information that would be helpful in implementing |
| 1547 | the task. |
| 1548 | |
| 1549 | Reply with ONLY the reprompt text. |
| 1550 | ` |
| 1551 | userMessage := ant.Message{ |
| 1552 | Role: "user", |
| 1553 | Content: []ant.Content{{Type: "text", Text: msg}}, |
| 1554 | } |
| 1555 | // By doing this in a subconversation, the agent doesn't call tools (because |
| 1556 | // there aren't any), and there's not a concurrency risk with on-going other |
| 1557 | // outstanding conversations. |
| 1558 | convo := a.convo.SubConvoWithHistory() |
| 1559 | resp, err := convo.SendMessage(userMessage) |
| 1560 | if err != nil { |
| 1561 | a.pushToOutbox(ctx, errorMessage(err)) |
| 1562 | return "", err |
| 1563 | } |
| 1564 | textContent := collectTextContent(resp) |
| 1565 | return textContent, nil |
| 1566 | } |