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