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

//go:build !tinygo && !wasm

// Package starlark provides a Starlark-based implementation of the CodeMode interface.
// Starlark is a Python-like language designed for configuration and embedded scripting.
// See https://github.com/google/starlark-go for more information.
package starlark

import (
	"fmt"
	"sync"
	"sync/atomic"
	"time"

	"github.com/maximhq/bifrost/core/mcp"
	"github.com/maximhq/bifrost/core/schemas"
)

// StarlarkCodeMode implements the CodeMode interface using a Starlark interpreter.
// It provides a sandboxed Python-like execution environment with access to MCP tools.
type StarlarkCodeMode struct {
	// Configuration (atomic for thread-safe updates)
	bindingLevel         atomic.Value // schemas.CodeModeBindingLevel
	toolExecutionTimeout atomic.Value // time.Duration

	// Dependencies
	clientManager          mcp.ClientManager
	pluginPipelineProvider func() mcp.PluginPipeline
	releasePluginPipeline  func(pipeline mcp.PluginPipeline)
	fetchNewRequestIDFunc  func(ctx *schemas.BifrostContext) string
	oauth2Provider         schemas.OAuth2Provider

	// Logger for this instance
	logger schemas.Logger

	// Mutex for protecting logs during concurrent execution
	logMu sync.Mutex
}

// NewStarlarkCodeMode creates a new Starlark-based CodeMode implementation.
//
// Parameters:
//   - config: Configuration for the code mode (binding level, timeouts). Can be nil for defaults.
//   - logger: Logger instance for this code mode. Can be nil.
//
// Returns:
//   - *StarlarkCodeMode: A new Starlark code mode instance
//
// Note: Dependencies must be set via SetDependencies before the CodeMode can execute tools.
// This allows the CodeMode to be created before the MCPManager, avoiding circular dependencies.
func NewStarlarkCodeMode(config *mcp.CodeModeConfig, logger schemas.Logger) *StarlarkCodeMode {
	if config == nil {
		config = mcp.DefaultCodeModeConfig()
	}

	if config.BindingLevel == "" {
		config.BindingLevel = schemas.CodeModeBindingLevelServer
	}

	if config.ToolExecutionTimeout <= 0 {
		config.ToolExecutionTimeout = schemas.DefaultToolExecutionTimeout
	}

	if logger == nil {
		logger = defaultLogger
	}

	s := &StarlarkCodeMode{
		logger: logger,
	}

	// Initialize atomic values
	s.bindingLevel.Store(config.BindingLevel)
	s.toolExecutionTimeout.Store(config.ToolExecutionTimeout)

	s.logger.Info("%s Starlark code mode initialized with binding level: %s, timeout: %v",
		mcp.CodeModeLogPrefix, config.BindingLevel, config.ToolExecutionTimeout)

	return s
}

// SetDependencies sets the dependencies required for code execution.
// This must be called after the MCPManager is created, as the dependencies
// include the ClientManager (which is the MCPManager itself).
func (s *StarlarkCodeMode) SetDependencies(deps *mcp.CodeModeDependencies) {
	if deps != nil {
		s.clientManager = deps.ClientManager
		s.pluginPipelineProvider = deps.PluginPipelineProvider
		s.releasePluginPipeline = deps.ReleasePluginPipeline
		s.fetchNewRequestIDFunc = deps.FetchNewRequestIDFunc
		s.oauth2Provider = deps.OAuth2Provider
	}
}

// GetTools returns the code mode meta-tools for Starlark execution.
// These tools allow LLMs to discover, read, and execute code against MCP servers.
func (s *StarlarkCodeMode) GetTools() []schemas.ChatTool {
	return []schemas.ChatTool{
		s.createListToolFilesTool(),
		s.createReadToolFileTool(),
		s.createGetToolDocsTool(),
		s.createExecuteToolCodeTool(),
	}
}

// ExecuteTool handles a code mode tool call.
// It dispatches to the appropriate handler based on the tool name.
//
// Parameters:
//   - ctx: Context for tool execution
//   - toolCall: The tool call to execute
//
// Returns:
//   - *schemas.ChatMessage: The tool response message
//   - error: Any error that occurred during execution
func (s *StarlarkCodeMode) ExecuteTool(ctx *schemas.BifrostContext, toolCall schemas.ChatAssistantMessageToolCall) (*schemas.ChatMessage, error) {
	if toolCall.Function.Name == nil {
		return nil, fmt.Errorf("tool call missing function name")
	}

	toolName := *toolCall.Function.Name

	switch toolName {
	case mcp.ToolTypeListToolFiles:
		return s.handleListToolFiles(ctx, toolCall)
	case mcp.ToolTypeReadToolFile:
		return s.handleReadToolFile(ctx, toolCall)
	case mcp.ToolTypeGetToolDocs:
		return s.handleGetToolDocs(ctx, toolCall)
	case mcp.ToolTypeExecuteToolCode:
		return s.handleExecuteToolCode(ctx, toolCall)
	default:
		return nil, fmt.Errorf("unknown code mode tool: %s", toolName)
	}
}

// IsCodeModeTool returns true if the given tool name is a code mode tool.
func (s *StarlarkCodeMode) IsCodeModeTool(toolName string) bool {
	return mcp.IsCodeModeTool(toolName)
}

// GetBindingLevel returns the current code mode binding level.
func (s *StarlarkCodeMode) GetBindingLevel() schemas.CodeModeBindingLevel {
	val := s.bindingLevel.Load()
	if val == nil {
		return schemas.CodeModeBindingLevelServer
	}
	return val.(schemas.CodeModeBindingLevel)
}

// UpdateConfig updates the code mode configuration atomically.
func (s *StarlarkCodeMode) UpdateConfig(config *mcp.CodeModeConfig) {
	if config == nil {
		return
	}

	if config.BindingLevel != "" {
		s.bindingLevel.Store(config.BindingLevel)
	}

	if config.ToolExecutionTimeout > 0 {
		s.toolExecutionTimeout.Store(config.ToolExecutionTimeout)
	}

	s.logger.Info("%s Starlark code mode configuration updated: binding level=%s, timeout=%v",
		mcp.CodeModeLogPrefix, config.BindingLevel, config.ToolExecutionTimeout)
}

// getToolExecutionTimeout returns the current tool execution timeout.
func (s *StarlarkCodeMode) getToolExecutionTimeout() time.Duration {
	val := s.toolExecutionTimeout.Load()
	if val == nil {
		return schemas.DefaultToolExecutionTimeout
	}
	return val.(time.Duration)
}