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