bifrost/core/mcp/codemode/starlark/executecode.go

//go:build !tinygo && !wasm

package starlark

import (
	"context"
	"fmt"
	"strings"
	"time"

	"github.com/bytedance/sonic"
	"github.com/mark3labs/mcp-go/mcp"

	codemcp "github.com/maximhq/bifrost/core/mcp"
	"github.com/maximhq/bifrost/core/mcp/utils"
	"github.com/maximhq/bifrost/core/schemas"
	"go.starlark.net/starlark"
	"go.starlark.net/starlarkstruct"
	"go.starlark.net/syntax"
)

// ExecutionResult represents the result of code execution
type ExecutionResult struct {
	Result      interface{}          `json:"result"`
	Logs        []string             `json:"logs"`
	Errors      *ExecutionError      `json:"errors,omitempty"`
	Environment ExecutionEnvironment `json:"environment"`
}

// ExecutionErrorType represents the type of execution error
type ExecutionErrorType string

const (
	ExecutionErrorTypeCompile ExecutionErrorType = "compile"
	ExecutionErrorTypeSyntax  ExecutionErrorType = "syntax"
	ExecutionErrorTypeRuntime ExecutionErrorType = "runtime"
)

// ExecutionError represents an error during code execution
type ExecutionError struct {
	Kind    ExecutionErrorType `json:"kind"` // "compile", "syntax", or "runtime"
	Message string             `json:"message"`
	Hints   []string           `json:"hints"`
}

// ExecutionEnvironment contains information about the execution environment
type ExecutionEnvironment struct {
	ServerKeys []string `json:"serverKeys"`
}

// createExecuteToolCodeTool creates the executeToolCode tool definition for code mode.
// This tool allows executing Python (Starlark) code in a sandboxed interpreter with access to MCP server tools.
func (s *StarlarkCodeMode) createExecuteToolCodeTool() schemas.ChatTool {
	executeToolCodeProps := schemas.NewOrderedMapFromPairs(
		schemas.KV("code", map[string]interface{}{
			"type": "string",
			"description": "Python (Starlark) code to execute. Tool calls are synchronous: result = server.tool(param=\"value\"). " +
				"Use print() for logging. Assign to 'result' variable to return a value. " +
				"Retry after fixing syntax or logic errors, especially for read-only flows. Before rerunning code that already made tool calls, inspect prior outputs and avoid replaying stateful operations. " +
				"Example: items = server.list_items()\nfor item in items:\n    print(item[\"name\"])\nresult = items",
		}),
	)
	return schemas.ChatTool{
		Type: schemas.ChatToolTypeFunction,
		Function: &schemas.ChatToolFunction{
			Name: codemcp.ToolTypeExecuteToolCode,
			Description: schemas.Ptr(
				"Executes Python code in a sandboxed Starlark interpreter with MCP server tool access. " +
					"Servers are exposed as global objects: result = serverName.toolName(param=\"value\"). " +
					"This is the final step of the four-tool code mode workflow: listToolFiles -> readToolFile -> (optional) getToolDocs -> executeToolCode. " +
					"If you have not already read a tool's .pyi stub in this conversation, do that before writing code. " +
					"Do NOT guess callable tool names from natural language or stale assumptions; use the exact identifier returned by listToolFiles/readToolFile. " +

					"STARLARK DIFFERENCES FROM PYTHON — READ BEFORE WRITING CODE: " +
					"1. NO try/except/finally/raise — error handling is not supported, and tool failures cannot be caught inside Starlark. " +
					"2. NO classes — use dicts and functions. " +
					"3. NO imports, direct network access, or direct filesystem access — use MCP tools instead. " +
					"4. NO is operator — use == for comparison. " +
					"5. NO f-strings — use % formatting: \"Hello %s, count=%d\" % (name, n). " +
					"6. Each executeToolCode call runs in a FRESH ISOLATED SCOPE — no variables, functions, or state persist between calls. Re-fetch data or store it via MCP tools (e.g., SQLite, FileSystem) if needed across calls. " +

					"SYNTAX NOTES: " +
					"• Synchronous calls — NO async/await: result = server.tool(arg=\"value\") " +
					"• Use keyword arguments: server.tool(param=\"value\") NOT server.tool({\"param\": \"value\"}) " +
					"• Access dict values with brackets: result[\"key\"] NOT result.key " +
					"• Use print() for logging/debugging " +
					"• List comprehensions: [x for x in items if x[\"active\"]] " +
					"• String escapes work normally: \"line1\\nline2\" produces a newline " +
					"• Triple-quoted strings for multiline: \"\"\"multi\\nline\"\"\" " +
					"• chr(10) for newline character, chr(9) for tab " +
					"• To return a value, assign to 'result': result = computed_value " +
					"• MCP tool calls are timeout-limited; avoid long or infinite loops " +

					"AVAILABLE BUILTINS: print, len, range, enumerate, zip, sorted, reversed, min, max, " +
					"int, float, str, bool, list, dict, tuple, set, hasattr, getattr, type, chr, ord, any, all, hash, repr. " +

					"RETRY POLICY: Retry after fixing syntax or logic errors, especially for read-only flows. Before rerunning code that already made tool calls, inspect prior outputs and avoid replaying stateful operations.",
			),

			Parameters: &schemas.ToolFunctionParameters{
				Type:       "object",
				Properties: executeToolCodeProps,
				Required:   []string{"code"},
			},
		},
	}
}

// handleExecuteToolCode handles the executeToolCode tool call.
func (s *StarlarkCodeMode) handleExecuteToolCode(ctx *schemas.BifrostContext, toolCall schemas.ChatAssistantMessageToolCall) (*schemas.ChatMessage, error) {
	toolName := "unknown"
	if toolCall.Function.Name != nil {
		toolName = *toolCall.Function.Name
	}
	s.logger.Debug("%s Handling executeToolCode tool call: %s", codemcp.CodeModeLogPrefix, toolName)

	// Parse tool arguments
	var arguments map[string]interface{}
	if err := sonic.Unmarshal([]byte(toolCall.Function.Arguments), &arguments); err != nil {
		s.logger.Debug("%s Failed to parse tool arguments: %v", codemcp.CodeModeLogPrefix, err)
		return nil, fmt.Errorf("failed to parse tool arguments: %v", err)
	}

	code, ok := arguments["code"].(string)
	if !ok || code == "" {
		s.logger.Debug("%s Code parameter missing or empty", codemcp.CodeModeLogPrefix)
		return nil, fmt.Errorf("code parameter is required and must be a non-empty string")
	}

	s.logger.Debug("%s Starting code execution", codemcp.CodeModeLogPrefix)
	result := s.executeCode(ctx, code)
	s.logger.Debug("%s Code execution completed. Success: %v, Has errors: %v, Log count: %d", codemcp.CodeModeLogPrefix, result.Errors == nil, result.Errors != nil, len(result.Logs))

	// Format response text
	var responseText string
	var executionSuccess bool = true
	if result.Errors != nil {
		s.logger.Debug("%s Formatting error response. Error kind: %s, Message length: %d, Hints count: %d", codemcp.CodeModeLogPrefix, result.Errors.Kind, len(result.Errors.Message), len(result.Errors.Hints))
		logsText := ""
		if len(result.Logs) > 0 {
			logsText = fmt.Sprintf("\n\nPrint Output:\n%s\n", strings.Join(result.Logs, "\n"))
		}

		responseText = fmt.Sprintf(
			"Execution %s error:\n\n%s\n\nHints:\n%s%s\n\nEnvironment:\n  Available server keys: %s",
			result.Errors.Kind,
			result.Errors.Message,
			strings.Join(result.Errors.Hints, "\n"),
			logsText,
			strings.Join(result.Environment.ServerKeys, ", "),
		)
		s.logger.Debug("%s Error response formatted. Response length: %d chars", codemcp.CodeModeLogPrefix, len(responseText))
	} else {
		hasLogs := len(result.Logs) > 0
		hasResult := result.Result != nil
		s.logger.Debug("%s Formatting success response. Has logs: %v, Has result: %v", codemcp.CodeModeLogPrefix, hasLogs, hasResult)

		if !hasLogs && !hasResult {
			executionSuccess = false
			s.logger.Debug("%s Execution completed with no data (no logs, no result), marking as failure", codemcp.CodeModeLogPrefix)
			hints := []string{
				"Add print() statements throughout your code to debug and see what's happening at each step",
				"Assign the final value to 'result' variable if you want to return it: result = computed_value",
				"Check that your tool calls are actually executing and returning data",
			}
			responseText = fmt.Sprintf(
				"Execution completed but produced no data:\n\n"+
					"The code executed without errors but returned no output (no print output and no result variable).\n\n"+
					"Hints:\n%s\n\n"+
					"Environment:\n  Available server keys: %s",
				strings.Join(hints, "\n"),
				strings.Join(result.Environment.ServerKeys, ", "),
			)
			s.logger.Debug("%s No-data failure response formatted. Response length: %d chars", codemcp.CodeModeLogPrefix, len(responseText))
		} else {
			if hasLogs {
				responseText = fmt.Sprintf("Print output:\n%s\n\nExecution completed successfully.",
					strings.Join(result.Logs, "\n"))
			} else {
				responseText = "Execution completed successfully."
			}
			if hasResult {
				resultJSON, err := schemas.MarshalSortedIndent(result.Result, "", "  ")
				if err == nil {
					responseText += fmt.Sprintf("\nReturn value: %s", string(resultJSON))
					s.logger.Debug("%s Added return value to response (JSON length: %d chars)", codemcp.CodeModeLogPrefix, len(resultJSON))
				} else {
					s.logger.Debug("%s Failed to marshal result to JSON: %v", codemcp.CodeModeLogPrefix, err)
				}
			}

			responseText += fmt.Sprintf("\n\nEnvironment:\n  Available server keys: %s",
				strings.Join(result.Environment.ServerKeys, ", "))
			responseText += "\nNote: This is a Starlark (Python subset) environment. Use MCP tools for external interactions."
			s.logger.Debug("%s Success response formatted. Response length: %d chars, Server keys: %v", codemcp.CodeModeLogPrefix, len(responseText), result.Environment.ServerKeys)
		}
	}

	s.logger.Debug("%s Returning tool response message. Execution success: %v", codemcp.CodeModeLogPrefix, executionSuccess)
	return createToolResponseMessage(toolCall, responseText), nil
}

// executeCode executes Python (Starlark) code in a sandboxed interpreter with MCP tool bindings.
func (s *StarlarkCodeMode) executeCode(ctx *schemas.BifrostContext, code string) ExecutionResult {
	logs := []string{}

	s.logger.Debug("%s Starting Starlark code execution", codemcp.CodeModeLogPrefix)

	// Step 1: Handle empty code
	trimmedCode := strings.TrimSpace(code)
	if trimmedCode == "" {
		return ExecutionResult{
			Result: nil,
			Logs:   logs,
			Errors: nil,
			Environment: ExecutionEnvironment{
				ServerKeys: []string{},
			},
		}
	}

	// Step 2: Build tool bindings for all connected servers
	availableToolsPerClient := s.clientManager.GetToolPerClient(ctx)
	serverKeys := make([]string, 0, len(availableToolsPerClient))
	predeclared := starlark.StringDict{}

	// Thread-safe log appender
	appendLog := func(msg string) {
		s.logMu.Lock()
		defer s.logMu.Unlock()
		logs = append(logs, msg)
	}

	s.logger.Debug("%s GetToolPerClient returned %d clients", codemcp.CodeModeLogPrefix, len(availableToolsPerClient))

	for clientName, tools := range availableToolsPerClient {
		client := s.clientManager.GetClientByName(clientName)
		if client == nil {
			s.logger.Warn("%s Client %s not found, skipping", codemcp.CodeModeLogPrefix, clientName)
			continue
		}
		s.logger.Debug("%s [%s] Client found. IsCodeModeClient: %v, ToolCount: %d", codemcp.CodeModeLogPrefix, clientName, client.ExecutionConfig.IsCodeModeClient, len(tools))
		if !client.ExecutionConfig.IsCodeModeClient || len(tools) == 0 {
			s.logger.Debug("%s [%s] Skipped: IsCodeModeClient=%v, HasTools=%v", codemcp.CodeModeLogPrefix, clientName, client.ExecutionConfig.IsCodeModeClient, len(tools) > 0)
			continue
		}
		serverKeys = append(serverKeys, clientName)

		// Build struct with tool methods
		structMembers := starlark.StringDict{}

		for _, tool := range tools {
			if tool.Function == nil || tool.Function.Name == "" {
				continue
			}

			originalToolName := tool.Function.Name
			parsedToolName := getCanonicalToolName(clientName, originalToolName)
			compatibilityAlias := getCompatibilityToolAlias(clientName, originalToolName)

			s.logger.Debug("%s [%s] Binding tool: %s -> %s", codemcp.CodeModeLogPrefix, clientName, originalToolName, parsedToolName)

			// Capture variables for closure
			capturedToolName := originalToolName
			capturedClientName := clientName

			// Create a Starlark builtin function for this tool
			toolFunc := starlark.NewBuiltin(parsedToolName, func(thread *starlark.Thread, fn *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
				// Convert kwargs to Go map
				goArgs := make(map[string]interface{})
				for _, kwarg := range kwargs {
					if len(kwarg) == 2 {
						key := string(kwarg[0].(starlark.String))
						value := starlarkToGo(kwarg[1])
						goArgs[key] = value
					}
				}

				// Also handle positional args if there's exactly one dict argument
				if len(args) == 1 && len(kwargs) == 0 {
					if dict, ok := args[0].(*starlark.Dict); ok {
						for _, item := range dict.Items() {
							if keyStr, ok := item[0].(starlark.String); ok {
								goArgs[string(keyStr)] = starlarkToGo(item[1])
							}
						}
					}
				}

				// Call the MCP tool
				result, err := s.callMCPTool(ctx, capturedClientName, capturedToolName, goArgs, appendLog)
				if err != nil {
					return starlark.None, fmt.Errorf("tool call failed: %v", err)
				}

				// Convert result back to Starlark
				return goToStarlark(result), nil
			})

			structMembers[parsedToolName] = toolFunc

			if compatibilityAlias != parsedToolName && isValidStarlarkIdentifier(compatibilityAlias) {
				if _, exists := structMembers[compatibilityAlias]; !exists {
					structMembers[compatibilityAlias] = toolFunc
					s.logger.Debug("%s [%s] Added compatibility alias: %s -> %s", codemcp.CodeModeLogPrefix, clientName, compatibilityAlias, parsedToolName)
				}
			}
		}

		// Create a struct for this server
		serverStruct := starlarkstruct.FromStringDict(starlark.String(clientName), structMembers)
		predeclared[clientName] = serverStruct
		s.logger.Debug("%s [%s] Added server struct with %d tools", codemcp.CodeModeLogPrefix, clientName, len(structMembers))
	}

	if len(serverKeys) > 0 {
		s.logger.Debug("%s Bound %d servers with tools: %v", codemcp.CodeModeLogPrefix, len(serverKeys), serverKeys)
	} else {
		s.logger.Debug("%s No servers available for code mode execution", codemcp.CodeModeLogPrefix)
	}

	// Step 3: Create Starlark thread with print function and timeout
	toolExecutionTimeout := s.getToolExecutionTimeout()
	timeoutCtx, cancel := context.WithTimeout(ctx, toolExecutionTimeout)
	defer cancel()

	thread := &starlark.Thread{
		Name: "codemode",
		Print: func(_ *starlark.Thread, msg string) {
			appendLog(msg)
		},
	}

	// Set up cancellation check — watch the context and cancel the Starlark
	// thread so that infinite loops and other long-running scripts are interrupted
	// when the execution timeout fires.
	thread.SetLocal("context", timeoutCtx)
	go func() {
		<-timeoutCtx.Done()
		thread.Cancel(timeoutCtx.Err().Error())
	}()

	// Step 4: Configure Starlark dialect options for a Python-like experience
	starlarkOpts := &syntax.FileOptions{
		TopLevelControl: true, // allow if/for/while at top level (not just inside functions)
		While:           true, // enable while loops
		Set:             true, // enable set() builtin
		GlobalReassign:  true, // allow reassignment to top-level names
		Recursion:       true, // allow recursive functions
	}

	// Step 5: Execute the code
	globals, err := starlark.ExecFileOptions(starlarkOpts, thread, "code.star", trimmedCode, predeclared)

	if err != nil {
		errorMessage := err.Error()
		hints := generatePythonErrorHints(errorMessage, serverKeys)
		s.logger.Debug("%s Execution failed: %s", codemcp.CodeModeLogPrefix, errorMessage)

		errorKind := ExecutionErrorTypeRuntime
		if strings.Contains(errorMessage, "syntax error") {
			errorKind = ExecutionErrorTypeSyntax
		}

		return ExecutionResult{
			Result: nil,
			Logs:   logs,
			Errors: &ExecutionError{
				Kind:    errorKind,
				Message: errorMessage,
				Hints:   hints,
			},
			Environment: ExecutionEnvironment{
				ServerKeys: serverKeys,
			},
		}
	}

	// Step 6: Extract result from globals
	var result interface{}
	if resultVal, ok := globals["result"]; ok && resultVal != starlark.None {
		result = starlarkToGo(resultVal)
	}

	s.logger.Debug("%s Execution completed successfully", codemcp.CodeModeLogPrefix)
	return ExecutionResult{
		Result: result,
		Logs:   logs,
		Errors: nil,
		Environment: ExecutionEnvironment{
			ServerKeys: serverKeys,
		},
	}
}

// callMCPTool calls an MCP tool and returns the result.
func (s *StarlarkCodeMode) callMCPTool(ctx *schemas.BifrostContext, clientName, toolName string, args map[string]interface{}, appendLog func(string)) (interface{}, error) {
	// Get available tools per client
	availableToolsPerClient := s.clientManager.GetToolPerClient(ctx)

	// Find the client by name
	tools, exists := availableToolsPerClient[clientName]
	if !exists || len(tools) == 0 {
		return nil, fmt.Errorf("client not found for server name: %s", clientName)
	}

	// Get client using a tool from this client
	var client *schemas.MCPClientState
	for _, tool := range tools {
		if tool.Function != nil && tool.Function.Name != "" {
			client = s.clientManager.GetClientForTool(tool.Function.Name)
			if client != nil {
				break
			}
		}
	}

	if client == nil {
		return nil, fmt.Errorf("client not found for server name: %s", clientName)
	}

	// Strip the client name prefix from tool name before calling MCP server
	originalToolName := stripClientPrefix(toolName, clientName)

	originalRequestID, ok := ctx.Value(schemas.BifrostContextKeyRequestID).(string)
	if !ok {
		originalRequestID = ""
	}

	// Generate new request ID for this nested tool call
	var newRequestID string
	if s.fetchNewRequestIDFunc != nil {
		newRequestID = s.fetchNewRequestIDFunc(ctx)
	} else {
		newRequestID = fmt.Sprintf("exec_%d_%s", time.Now().UnixNano(), toolName)
	}

	// Create new child context
	deadline, hasDeadline := ctx.Deadline()
	if !hasDeadline {
		deadline = schemas.NoDeadline
	}
	nestedCtx := schemas.NewBifrostContext(ctx, deadline)
	nestedCtx.SetValue(schemas.BifrostContextKeyRequestID, newRequestID)
	if originalRequestID != "" {
		nestedCtx.SetValue(schemas.BifrostContextKeyParentMCPRequestID, originalRequestID)
	}

	// Marshal arguments to JSON for the tool call
	argsJSON, err := schemas.MarshalSorted(args)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal tool arguments: %v", err)
	}

	// Build tool call for MCP request
	toolCallReq := schemas.ChatAssistantMessageToolCall{
		ID: schemas.Ptr(newRequestID),
		Function: schemas.ChatAssistantMessageToolCallFunction{
			Name:      schemas.Ptr(toolName),
			Arguments: string(argsJSON),
		},
	}

	// Create BifrostMCPRequest
	mcpRequest := &schemas.BifrostMCPRequest{
		RequestType:                  schemas.MCPRequestTypeChatToolCall,
		ChatAssistantMessageToolCall: &toolCallReq,
	}

	// Check if plugin pipeline is available
	if s.pluginPipelineProvider == nil {
		// Should never happen, but just in case
		s.logger.Warn("%s Plugin pipeline provider is nil", codemcp.CodeModeLogPrefix)
		return nil, fmt.Errorf("plugin pipeline provider is nil")
	}

	// Get plugin pipeline and run hooks
	pipeline := s.pluginPipelineProvider()
	if pipeline == nil {
		// Should never happen, but just in case
		s.logger.Warn("%s Plugin pipeline is nil", codemcp.CodeModeLogPrefix)
		return nil, fmt.Errorf("plugin pipeline is nil")
	}
	defer s.releasePluginPipeline(pipeline)

	// Run PreMCPHooks
	preReq, shortCircuit, preCount := pipeline.RunMCPPreHooks(nestedCtx, mcpRequest)

	// Handle short-circuit cases
	if shortCircuit != nil {
		if shortCircuit.Response != nil {
			finalResp, _ := pipeline.RunMCPPostHooks(nestedCtx, shortCircuit.Response, nil, preCount)
			if finalResp != nil {
				if finalResp.ChatMessage != nil {
					return extractResultFromChatMessage(finalResp.ChatMessage), nil
				}
				if finalResp.ResponsesMessage != nil {
					result, err := extractResultFromResponsesMessage(finalResp.ResponsesMessage)
					if err != nil {
						return nil, err
					}
					if result != nil {
						return result, nil
					}
				}
			}
			return nil, fmt.Errorf("plugin short-circuit returned invalid response")
		}
		if shortCircuit.Error != nil {
			pipeline.RunMCPPostHooks(nestedCtx, nil, shortCircuit.Error, preCount)
			if shortCircuit.Error.Error != nil {
				return nil, fmt.Errorf("%s", shortCircuit.Error.Error.Message)
			}
			return nil, fmt.Errorf("plugin short-circuit error")
		}
	}

	// If pre-hooks modified the request, extract updated args
	if preReq != nil && preReq.ChatAssistantMessageToolCall != nil {
		toolCallReq = *preReq.ChatAssistantMessageToolCall
		if toolCallReq.Function.Arguments != "" {
			if err := sonic.Unmarshal([]byte(toolCallReq.Function.Arguments), &args); err != nil {
				s.logger.Warn("%s Failed to parse modified tool arguments, using original: %v", codemcp.CodeModeLogPrefix, err)
			}
		}
	}

	// Execute tool
	startTime := time.Now()
	toolNameToCall := originalToolName

	callRequest := mcp.CallToolRequest{
		Request: mcp.Request{
			Method: string(mcp.MethodToolsCall),
		},
		Params: mcp.CallToolParams{
			Name:      toolNameToCall,
			Arguments: args,
		},
		Header: utils.GetHeadersForToolExecution(nestedCtx, client),
	}

	toolExecutionTimeout := s.getToolExecutionTimeout()
	toolCtx, cancel := context.WithTimeout(nestedCtx, toolExecutionTimeout)
	defer cancel()

	var toolResponse *mcp.CallToolResult
	var callErr error

	if client.ExecutionConfig.AuthType == schemas.MCPAuthTypePerUserOauth {
		accessToken, err := utils.ResolvePerUserOAuthToken(nestedCtx, client, s.oauth2Provider)
		if err != nil {
			return nil, err
		}

		if client.Conn == nil {
			// Per-user OAuth with no persistent connection — use a temporary connection.
			// Assign to outer toolResponse/callErr so the shared logging + post-hooks path runs.
			toolResponse, callErr = codemcp.ExecuteToolWithUserToken(toolCtx, client.ExecutionConfig, toolNameToCall, args, accessToken, s.logger)
			if callErr != nil && toolCtx.Err() == context.DeadlineExceeded {
				callErr = fmt.Errorf("MCP tool call timed out after %v: %s", toolExecutionTimeout, toolName)
			}
		} else {
			callRequest.Header = utils.BuildPerUserOAuthHeaders(callRequest.Header, accessToken)
			toolResponse, callErr = client.Conn.CallTool(toolCtx, callRequest)
		}
	} else {
		toolResponse, callErr = client.Conn.CallTool(toolCtx, callRequest)
	}

	latency := time.Since(startTime).Milliseconds()

	var mcpResp *schemas.BifrostMCPResponse
	var bifrostErr *schemas.BifrostError

	if callErr != nil {
		s.logger.Debug("%s Tool call failed: %s.%s - %v", codemcp.CodeModeLogPrefix, clientName, toolName, callErr)
		appendLog(fmt.Sprintf("[TOOL] %s.%s error: %v", clientName, toolName, callErr))
		bifrostErr = &schemas.BifrostError{
			IsBifrostError: false,
			Error: &schemas.ErrorField{
				Message: fmt.Sprintf("tool call failed for %s.%s: %v", clientName, toolName, callErr),
			},
		}
	} else {
		rawResult := extractTextFromMCPResponse(toolResponse, toolName)

		if after, ok := strings.CutPrefix(rawResult, "Error: "); ok {
			errorMsg := after
			s.logger.Debug("%s Tool returned error result: %s.%s - %s", codemcp.CodeModeLogPrefix, clientName, toolName, errorMsg)
			appendLog(fmt.Sprintf("[TOOL] %s.%s error result: %s", clientName, toolName, errorMsg))
			bifrostErr = &schemas.BifrostError{
				IsBifrostError: false,
				Error: &schemas.ErrorField{
					Message: errorMsg,
				},
			}
		} else {
			mcpResp = &schemas.BifrostMCPResponse{
				ChatMessage: createToolResponseMessage(toolCallReq, rawResult),
				ExtraFields: schemas.BifrostMCPResponseExtraFields{
					ClientName: clientName,
					ToolName:   originalToolName,
					Latency:    latency,
				},
			}

			resultStr := formatResultForLog(rawResult)
			logToolName := stripClientPrefix(toolName, clientName)
			logToolName = strings.ReplaceAll(logToolName, "-", "_")
			appendLog(fmt.Sprintf("[TOOL] %s.%s raw response: %s", clientName, logToolName, resultStr))
		}
	}

	// Run post-hooks
	finalResp, finalErr := pipeline.RunMCPPostHooks(nestedCtx, mcpResp, bifrostErr, preCount)

	if finalErr != nil {
		if finalErr.Error != nil {
			return nil, fmt.Errorf("%s", finalErr.Error.Message)
		}
		return nil, fmt.Errorf("tool execution failed")
	}

	if finalResp == nil {
		return nil, fmt.Errorf("plugin post-hooks returned invalid response")
	}

	if finalResp.ChatMessage != nil {
		return extractResultFromChatMessage(finalResp.ChatMessage), nil
	}

	if finalResp.ResponsesMessage != nil {
		result, err := extractResultFromResponsesMessage(finalResp.ResponsesMessage)
		if err != nil {
			return nil, err
		}
		if result != nil {
			return result, nil
		}
	}

	return nil, fmt.Errorf("plugin post-hooks returned invalid response")
}