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