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