Skip to content

Commit 311e34c

Browse files
ysyneuysyneu
committed
feat: enhance logging and observability with trace IDs and structured logs
- Added TraceID to TaskRequestPayload for end-to-end request tracking - Implemented OrderedTextHandler for more consistent log output - Updated MCP call and task execution flows to use context-aware structured logging - Improved error reporting in MCP tool calls
1 parent b303f1c commit 311e34c

5 files changed

Lines changed: 53 additions & 22 deletions

File tree

‎cmd/main.go‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ import (
1313

1414
"github.com/spf13/cobra"
1515

16+
"github.com/flashcatcloud/flashduty-runner/log"
1617
"github.com/flashcatcloud/flashduty-runner/permission"
1718
"github.com/flashcatcloud/flashduty-runner/workspace"
1819
"github.com/flashcatcloud/flashduty-runner/ws"
@@ -244,7 +245,7 @@ func runRunner() error {
244245
func setupLogging(levelStr string) {
245246
level := parseLogLevel(levelStr)
246247
opts := &slog.HandlerOptions{Level: level}
247-
handler := slog.NewTextHandler(os.Stdout, opts)
248+
handler := log.NewOrderedTextHandler(os.Stdout, opts)
248249
slog.SetDefault(slog.New(handler))
249250
}
250251

‎mcp/client.go‎

Lines changed: 23 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -110,30 +110,34 @@ func (m *ClientManager) GetSession(ctx context.Context, server *protocol.MCPServ
110110
}
111111

112112
// CallTool executes a single MCP tool call using the manager for session persistence.
113-
func (m *ClientManager) CallTool(ctx context.Context, server *protocol.MCPServerConfig, toolName string, arguments map[string]any) (*sdk_mcp.CallToolResult, error) {
114-
slog.Info("mcp call tool",
113+
func (m *ClientManager) CallTool(ctx context.Context, server *protocol.MCPServerConfig, toolName string, arguments map[string]any, logger *slog.Logger) (*sdk_mcp.CallToolResult, error) {
114+
if logger == nil {
115+
logger = slog.Default()
116+
}
117+
118+
logger.Info("mcp call tool",
115119
"server_name", server.Name,
116120
"tool_name", toolName,
117121
"arguments", arguments,
118122
)
119123

120124
session, err := m.GetSession(ctx, server)
121125
if err != nil {
122-
slog.Error("mcp get session failed", "server_name", server.Name, "error", err)
126+
logger.Error("mcp get session failed", "server_name", server.Name, "error", err)
123127
return nil, err
124128
}
125129

126130
callCtx, callCancel := context.WithTimeout(ctx, DefaultCallTimeout)
127131
defer callCancel()
128132

129-
slog.Debug("mcp calling tool", "server_name", server.Name, "tool_name", toolName, "timeout", DefaultCallTimeout)
133+
logger.Debug("mcp calling tool", "server_name", server.Name, "tool_name", toolName, "timeout", DefaultCallTimeout)
130134

131135
result, err := session.CallTool(callCtx, &sdk_mcp.CallToolParams{
132136
Name: toolName,
133137
Arguments: arguments,
134138
})
135139
if err != nil {
136-
slog.Error("mcp call tool failed",
140+
logger.Error("mcp call tool failed",
137141
"server_name", server.Name,
138142
"tool_name", toolName,
139143
"error", err,
@@ -142,11 +146,19 @@ func (m *ClientManager) CallTool(ctx context.Context, server *protocol.MCPServer
142146
return nil, fmt.Errorf("failed to call tool '%s': %w", toolName, err)
143147
}
144148

145-
slog.Info("mcp call tool success",
146-
"server_name", server.Name,
147-
"tool_name", toolName,
148-
"is_error", result.IsError,
149-
)
149+
if result.IsError {
150+
errorContent := ExtractContent(result)
151+
logger.Warn("mcp tool returned error",
152+
"server_name", server.Name,
153+
"tool_name", toolName,
154+
"error_content", errorContent,
155+
)
156+
} else {
157+
logger.Info("mcp call tool success",
158+
"server_name", server.Name,
159+
"tool_name", toolName,
160+
)
161+
}
150162

151163
return result, nil
152164
}
@@ -210,7 +222,7 @@ func GetDefaultManager() *ClientManager {
210222

211223
// CallTool executes a single MCP tool call using the default manager.
212224
func CallTool(ctx context.Context, server *protocol.MCPServerConfig, toolName string, arguments map[string]any) (*sdk_mcp.CallToolResult, error) {
213-
return GetDefaultManager().CallTool(ctx, server, toolName, arguments)
225+
return GetDefaultManager().CallTool(ctx, server, toolName, arguments, nil)
214226
}
215227

216228
// ListTools lists available tools from an MCP server using the default manager.

‎protocol/messages.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -116,6 +116,7 @@ const (
116116
// TaskRequestPayload is the payload for task request messages.
117117
type TaskRequestPayload struct {
118118
TaskID string `json:"task_id"`
119+
TraceID string `json:"trace_id,omitempty"`
119120
SourceInstanceID string `json:"source_instance_id,omitempty"` // Safari instance ID
120121
Operation TaskOperation `json:"operation"`
121122
Args json.RawMessage `json:"args"`

‎workspace/workspace.go‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import (
1010
"encoding/json"
1111
"fmt"
1212
"io"
13+
"log/slog"
1314
"os"
1415
"os/exec"
1516
"path/filepath"
@@ -482,7 +483,7 @@ func (w *Workspace) WriteRaw(ctx context.Context, path string, content []byte) e
482483
}
483484

484485
// MCPCall executes an MCP tool call.
485-
func (w *Workspace) MCPCall(ctx context.Context, args *protocol.MCPCallArgs) (*protocol.MCPCallResult, error) {
486+
func (w *Workspace) MCPCall(ctx context.Context, args *protocol.MCPCallArgs, logger *slog.Logger) (*protocol.MCPCallResult, error) {
486487
// Parse arguments
487488
var toolArgs map[string]any
488489
if len(args.Args) > 0 {
@@ -492,7 +493,7 @@ func (w *Workspace) MCPCall(ctx context.Context, args *protocol.MCPCallArgs) (*p
492493
}
493494

494495
// Call MCP tool
495-
result, err := w.mcpMgr.CallTool(ctx, &args.Server, args.ToolName, toolArgs)
496+
result, err := w.mcpMgr.CallTool(ctx, &args.Server, args.ToolName, toolArgs, logger)
496497
if err != nil {
497498
return nil, err
498499
}

‎ws/handler.go‎

Lines changed: 24 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -94,7 +94,12 @@ func (h *Handler) handleTaskRequest(ctx context.Context, msg *protocol.Message)
9494
return fmt.Errorf("failed to unmarshal task request: %w", err)
9595
}
9696

97-
slog.Info("received task request", "task_id", req.TaskID, "operation", req.Operation)
97+
logger := slog.With(
98+
"task_id", req.TaskID,
99+
"trace_id", req.TraceID,
100+
"operation", req.Operation,
101+
)
102+
logger.Info("received task request")
98103

99104
taskCtx, cancel := context.WithCancel(ctx)
100105
h.mu.Lock()
@@ -104,7 +109,7 @@ func (h *Handler) handleTaskRequest(ctx context.Context, msg *protocol.Message)
104109
h.taskWg.Add(1)
105110
go func() {
106111
defer h.taskWg.Done()
107-
h.executeAndSendResult(taskCtx, &req)
112+
h.executeAndSendResult(taskCtx, &req, logger)
108113
}()
109114
return nil
110115
}
@@ -117,15 +122,19 @@ func (h *Handler) unregisterTask(taskID string) {
117122
}
118123

119124
// executeAndSendResult executes a task and sends the result.
120-
func (h *Handler) executeAndSendResult(ctx context.Context, req *protocol.TaskRequestPayload) {
125+
func (h *Handler) executeAndSendResult(ctx context.Context, req *protocol.TaskRequestPayload, logger *slog.Logger) {
121126
defer h.unregisterTask(req.TaskID)
122127

123-
result, err := h.executeTask(ctx, req)
128+
logger.Info("executing task")
129+
130+
result, err := h.executeTask(ctx, req, logger)
124131
if err != nil {
132+
logger.Error("task execution failed", "error", err)
125133
h.sendTaskResult(req.TaskID, req.SourceInstanceID, false, nil, err.Error(), 1)
126134
return
127135
}
128136

137+
logger.Info("task completed successfully")
129138
h.sendTaskResult(req.TaskID, req.SourceInstanceID, true, result, "", 0)
130139
}
131140

@@ -137,7 +146,11 @@ func parseArgs[T any](data json.RawMessage) (*T, error) {
137146
return &args, nil
138147
}
139148

140-
func (h *Handler) executeTask(ctx context.Context, req *protocol.TaskRequestPayload) (any, error) {
149+
func (h *Handler) executeTask(ctx context.Context, req *protocol.TaskRequestPayload, logger *slog.Logger) (any, error) {
150+
if logger == nil {
151+
logger = slog.Default()
152+
}
153+
141154
switch req.Operation {
142155
case protocol.TaskOpRead:
143156
args, err := parseArgs[protocol.ReadArgs](req.Args)
@@ -196,15 +209,15 @@ func (h *Handler) executeTask(ctx context.Context, req *protocol.TaskRequestPayl
196209
if err != nil {
197210
return nil, fmt.Errorf("invalid mcp_call args: %w", err)
198211
}
199-
slog.Info("mcp_call task",
212+
logger.Info("mcp_call task",
200213
"server_name", args.Server.Name,
201214
"server_transport", args.Server.Transport,
202215
"server_url", args.Server.URL,
203216
"tool_name", args.ToolName,
204217
"has_headers", len(args.Server.Headers) > 0,
205218
"has_dynamic_headers", len(args.Server.DynamicHeaders) > 0,
206219
)
207-
return h.ws.MCPCall(ctx, args)
220+
return h.ws.MCPCall(ctx, args, logger)
208221

209222
case protocol.TaskOpMCPListTools:
210223
args, err := parseArgs[protocol.MCPListToolsArgs](req.Args)
@@ -278,12 +291,15 @@ func (h *Handler) handleMCPCall(ctx context.Context, msg *protocol.Message) erro
278291
return fmt.Errorf("failed to unmarshal mcp call: %w", err)
279292
}
280293

294+
// Create logger for MCP call (no task_id/trace_id available in this context)
295+
logger := slog.Default()
296+
281297
go func() {
282298
result, err := h.ws.MCPCall(ctx, &protocol.MCPCallArgs{
283299
Server: payload.Server,
284300
ToolName: payload.ToolName,
285301
Args: payload.Arguments,
286-
})
302+
}, logger)
287303
if err != nil {
288304
h.sendMCPResult(payload.CallID, false, nil, err.Error())
289305
} else {

0 commit comments

Comments
 (0)