Skip to content

Commit 311ed8a

Browse files
authored
feat(claude/tool_call): add tool call logging and improve format handling (#163)
1 parent 5e25714 commit 311ed8a

File tree

9 files changed

+164
-63
lines changed

9 files changed

+164
-63
lines changed

lib/httpapi/server.go

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -233,6 +233,10 @@ func NewServer(ctx context.Context, config ServerConfig) (*Server, error) {
233233
return mf.IsAgentReadyForInitialPrompt(config.AgentType, message)
234234
}
235235

236+
formatToolCall := func(message string) (string, []string) {
237+
return mf.FormatToolCall(config.AgentType, message)
238+
}
239+
236240
conversation := st.NewConversation(ctx, st.ConversationConfig{
237241
AgentType: config.AgentType,
238242
AgentIO: config.Process,
@@ -243,6 +247,8 @@ func NewServer(ctx context.Context, config ServerConfig) (*Server, error) {
243247
ScreenStabilityLength: 2 * time.Second,
244248
FormatMessage: formatMessage,
245249
ReadyForInitialPrompt: isAgentReadyForInitialPrompt,
250+
FormatToolCall: formatToolCall,
251+
Logger: logger,
246252
}, config.InitialPrompt)
247253
emitter := NewEventEmitter(1024)
248254

lib/msgfmt/format_tool_call.go

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
package msgfmt
2+
3+
import (
4+
"strings"
5+
)
6+
7+
func removeClaudeReportTaskToolCall(msg string) (string, []string) {
8+
msg = "\n" + msg // This handles the case where the message starts with a tool call
9+
10+
// Remove all tool calls that start with `● coder - coder_report_task (MCP)`
11+
lines := strings.Split(msg, "\n")
12+
13+
toolCallStartIdx := -1
14+
15+
// Store all tool call start and end indices [[start, end], ...]
16+
var toolCallIdxs [][]int
17+
18+
for i := 1; i < len(lines)-1; i++ {
19+
prevLine := strings.TrimSpace(lines[i-1])
20+
line := strings.TrimSpace(lines[i])
21+
nextLine := strings.TrimSpace(lines[i+1])
22+
23+
if strings.Contains(line, "coder - coder_report_task (MCP)") {
24+
toolCallStartIdx = i
25+
} else if toolCallStartIdx != -1 && line == "\"message\": \"Thanks for reporting!\"" && nextLine == "}" && strings.HasSuffix(prevLine, "{") {
26+
// Store [start, end] pair
27+
toolCallIdxs = append(toolCallIdxs, []int{toolCallStartIdx, min(len(lines), i+2)})
28+
29+
// Reset to find the next tool call
30+
toolCallStartIdx = -1
31+
}
32+
}
33+
34+
// If no tool calls found, return original message
35+
if len(toolCallIdxs) == 0 {
36+
return strings.TrimLeft(msg, "\n"), []string{}
37+
}
38+
39+
toolCallMessages := make([]string, 0)
40+
41+
// Remove tool calls from the message
42+
for i := len(toolCallIdxs) - 1; i >= 0; i-- {
43+
idxPair := toolCallIdxs[i]
44+
start, end := idxPair[0], idxPair[1]
45+
46+
toolCallMessages = append(toolCallMessages, strings.Join(lines[start:end], "\n"))
47+
48+
lines = append(lines[:start], lines[end:]...)
49+
}
50+
return strings.TrimLeft(strings.Join(lines, "\n"), "\n"), toolCallMessages
51+
}
52+
53+
func FormatToolCall(agentType AgentType, message string) (string, []string) {
54+
switch agentType {
55+
case AgentTypeClaude:
56+
return removeClaudeReportTaskToolCall(message)
57+
case AgentTypeGoose:
58+
return message, []string{}
59+
case AgentTypeAider:
60+
return message, []string{}
61+
case AgentTypeCodex:
62+
return message, []string{}
63+
case AgentTypeGemini:
64+
return message, []string{}
65+
case AgentTypeCopilot:
66+
return message, []string{}
67+
case AgentTypeAmp:
68+
return message, []string{}
69+
case AgentTypeCursor:
70+
return message, []string{}
71+
case AgentTypeAuggie:
72+
return message, []string{}
73+
case AgentTypeAmazonQ:
74+
return message, []string{}
75+
case AgentTypeOpencode:
76+
return message, []string{}
77+
case AgentTypeCustom:
78+
return message, []string{}
79+
default:
80+
return message, []string{}
81+
}
82+
}

lib/msgfmt/message_box.go

Lines changed: 0 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -100,53 +100,3 @@ func removeAmpMessageBox(msg string) string {
100100
}
101101
return formattedMsg
102102
}
103-
104-
func removeClaudeReportTaskToolCall(msg string) string {
105-
// Remove all tool calls that start with `● coder - coder_report_task (MCP)` till we encounter the next line starting with ●
106-
lines := strings.Split(msg, "\n")
107-
108-
toolCallStartIdx := -1
109-
newLineAfterToolCallIdx := -1
110-
111-
// Store all tool call start and end indices [[start, end], ...]
112-
var toolCallIdxs [][]int
113-
114-
for i := 0; i < len(lines); i++ {
115-
line := strings.TrimSpace(lines[i])
116-
117-
if strings.HasPrefix(line, "● coder - coder_report_task (MCP)") {
118-
toolCallStartIdx = i
119-
} else if toolCallStartIdx != -1 && strings.HasPrefix(line, "●") {
120-
// Store [start, end] pair
121-
toolCallIdxs = append(toolCallIdxs, []int{toolCallStartIdx, i})
122-
123-
// Reset to find the next tool call
124-
toolCallStartIdx = -1
125-
newLineAfterToolCallIdx = -1
126-
}
127-
if len(line) == 0 && toolCallStartIdx != -1 && newLineAfterToolCallIdx == -1 {
128-
newLineAfterToolCallIdx = i
129-
}
130-
}
131-
132-
// Handle the case where the last tool call goes till the end of the message
133-
// And a failsafe when the next message is not prefixed with ●
134-
if toolCallStartIdx != -1 && newLineAfterToolCallIdx != -1 {
135-
toolCallIdxs = append(toolCallIdxs, []int{toolCallStartIdx, newLineAfterToolCallIdx})
136-
}
137-
138-
// If no tool calls found, return original message
139-
if len(toolCallIdxs) == 0 {
140-
return msg
141-
}
142-
143-
// Remove tool calls from the message
144-
for i := len(toolCallIdxs) - 1; i >= 0; i-- {
145-
idxPair := toolCallIdxs[i]
146-
start, end := idxPair[0], idxPair[1]
147-
148-
lines = append(lines[:start], lines[end:]...)
149-
}
150-
151-
return strings.Join(lines, "\n")
152-
}

lib/msgfmt/msgfmt.go

Lines changed: 1 addition & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -254,14 +254,6 @@ func formatGenericMessage(message string, userInput string, agentType AgentType)
254254
return message
255255
}
256256

257-
func formatClaudeMessage(message string, userInput string) string {
258-
message = RemoveUserInput(message, userInput, AgentTypeClaude)
259-
message = removeMessageBox(message)
260-
message = removeClaudeReportTaskToolCall(message)
261-
message = trimEmptyLines(message)
262-
return message
263-
}
264-
265257
func formatCodexMessage(message string, userInput string) string {
266258
message = RemoveUserInput(message, userInput, AgentTypeCodex)
267259
message = removeCodexInputBox(message)
@@ -286,7 +278,7 @@ func formatAmpMessage(message string, userInput string) string {
286278
func FormatAgentMessage(agentType AgentType, message string, userInput string) string {
287279
switch agentType {
288280
case AgentTypeClaude:
289-
return formatClaudeMessage(message, userInput)
281+
return formatGenericMessage(message, userInput, agentType)
290282
case AgentTypeGoose:
291283
return formatGenericMessage(message, userInput, agentType)
292284
case AgentTypeAider:

lib/msgfmt/msgfmt_test.go

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -233,7 +233,16 @@ func TestFormatAgentMessage(t *testing.T) {
233233
assert.NoError(t, err)
234234
expected, err := testdataDir.ReadFile(path.Join(dir, string(agentType), c.Name(), "expected.txt"))
235235
assert.NoError(t, err)
236-
assert.Equal(t, string(expected), FormatAgentMessage(agentType, string(msg), string(userInput)))
236+
output, toolCalls := FormatToolCall(agentType, FormatAgentMessage(agentType, string(msg), string(userInput)))
237+
assert.Equal(t, string(expected), output)
238+
239+
// Assert on the tool calls if there's an expected file
240+
expectedToolCallsPath := path.Join(dir, string(agentType), c.Name(), "expected_tool_calls.txt")
241+
if expectedToolCallsData, err := testdataDir.ReadFile(expectedToolCallsPath); err == nil {
242+
expectedToolCalls := string(expectedToolCallsData)
243+
actualToolCalls := strings.Join(toolCalls, "\n---\n")
244+
assert.Equal(t, expectedToolCalls, actualToolCalls)
245+
}
237246
})
238247
}
239248
})

lib/msgfmt/testdata/format/claude/remove-task-tool-call/expected.txt

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
● I'll build a snake game for you. Let me start by reporting my
22
progress and creating a task list.
33

4+
45
● Now I'll create a complete snake game with HTML, CSS, and
56
JavaScript:
67

@@ -19,6 +20,8 @@
1920
padding: 0;
2021
… +334 lines (ctrl+o to expand)
2122

23+
24+
2225
● I've built a complete snake game for you! The game is saved
2326
at /home/coder/snake-game.html.
2427

@@ -36,4 +39,4 @@
3639
How to play:
3740
Open the HTML file in your web browser and use the arrow keys
3841
to move the snake. Collect the red food to grow and increase
39-
your score!
42+
your score!
Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
● coder - coder_report_task (MCP)(summary: "Snake game created
2+
successfully at snake-game.html",
3+
link: "file:///home/coder/snake-ga
4+
me.html", state: "working")
5+
⎿ {
6+
"message": "Thanks for reporting!"
7+
}
8+
---
9+
● coder - coder_report_task (MCP)(summary: "Snake game created
10+
successfully at snake-game.html",
11+
link: "file:///home/coder/snake-ga
12+
me.html", state: "working")
13+
⎿ {
14+
"message": "Thanks for reporting!"
15+
}
16+
---
17+
● coder - coder_report_task (MCP)(summary: "Building a snake game
18+
with HTML/CSS/JavaScript", link:
19+
"", state: "working")
20+
⎿ {
21+
"message": "Thanks for reporting!"
22+
}
23+
---
24+
● coder - coder_report_task (MCP)(summary: "Snake game created
25+
successfully at snake-game.html",
26+
link: "file:///home/coder/snake-ga
27+
me.html", state: "working")
28+
⎿ {
29+
"message": "Thanks for reporting!"
30+
}

lib/msgfmt/testdata/format/claude/remove-task-tool-call/msg.txt

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,13 @@
11
> Build a snake game
22

3+
● coder - coder_report_task (MCP)(summary: "Snake game created
4+
successfully at snake-game.html",
5+
link: "file:///home/coder/snake-ga
6+
me.html", state: "working")
7+
⎿ {
8+
"message": "Thanks for reporting!"
9+
}
10+
311
● I'll build a snake game for you. Let me start by reporting my
412
progress and creating a task list.
513

lib/screentracker/conversation.go

Lines changed: 23 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package screentracker
33
import (
44
"context"
55
"fmt"
6+
"log/slog"
67
"strings"
78
"sync"
89
"time"
@@ -43,6 +44,9 @@ type ConversationConfig struct {
4344
SkipSendMessageStatusCheck bool
4445
// ReadyForInitialPrompt detects whether the agent has initialized and is ready to accept the initial prompt
4546
ReadyForInitialPrompt func(message string) bool
47+
// FormatToolCall removes the coder report_task tool call from the agent message and also returns the array of removed tool calls
48+
FormatToolCall func(message string) (string, []string)
49+
Logger *slog.Logger
4650
}
4751

4852
type ConversationRole string
@@ -82,6 +86,8 @@ type Conversation struct {
8286
InitialPromptSent bool
8387
// ReadyForInitialPrompt keeps track if the agent is ready to accept the initial prompt
8488
ReadyForInitialPrompt bool
89+
// toolCallMessageSet keeps track of the tool calls that have been detected & logged in the current agent message
90+
toolCallMessageSet map[string]bool
8591
}
8692

8793
type ConversationStatus string
@@ -115,8 +121,9 @@ func NewConversation(ctx context.Context, cfg ConversationConfig, initialPrompt
115121
Time: cfg.GetTime(),
116122
},
117123
},
118-
InitialPrompt: initialPrompt,
119-
InitialPromptSent: len(initialPrompt) == 0,
124+
InitialPrompt: initialPrompt,
125+
InitialPromptSent: len(initialPrompt) == 0,
126+
toolCallMessageSet: make(map[string]bool),
120127
}
121128
return c
122129
}
@@ -205,9 +212,19 @@ func (c *Conversation) lastMessage(role ConversationRole) ConversationMessage {
205212
func (c *Conversation) updateLastAgentMessage(screen string, timestamp time.Time) {
206213
agentMessage := FindNewMessage(c.screenBeforeLastUserMessage, screen, c.cfg.AgentType)
207214
lastUserMessage := c.lastMessage(ConversationRoleUser)
215+
var toolCalls []string
208216
if c.cfg.FormatMessage != nil {
209217
agentMessage = c.cfg.FormatMessage(agentMessage, lastUserMessage.Message)
210218
}
219+
if c.cfg.FormatToolCall != nil {
220+
agentMessage, toolCalls = c.cfg.FormatToolCall(agentMessage)
221+
}
222+
for _, toolCall := range toolCalls {
223+
if c.toolCallMessageSet[toolCall] == false {
224+
c.toolCallMessageSet[toolCall] = true
225+
c.cfg.Logger.Info("Tool call detected", "toolCall", toolCall)
226+
}
227+
}
211228
shouldCreateNewMessage := len(c.messages) == 0 || c.messages[len(c.messages)-1].Role == ConversationRoleUser
212229
lastAgentMessage := c.lastMessage(ConversationRoleAgent)
213230
if lastAgentMessage.Message == agentMessage {
@@ -220,6 +237,10 @@ func (c *Conversation) updateLastAgentMessage(screen string, timestamp time.Time
220237
}
221238
if shouldCreateNewMessage {
222239
c.messages = append(c.messages, conversationMessage)
240+
241+
// Cleanup
242+
c.toolCallMessageSet = make(map[string]bool)
243+
223244
} else {
224245
c.messages[len(c.messages)-1] = conversationMessage
225246
}

0 commit comments

Comments
 (0)