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