goclaw/internal/tools/context_keys.go

package tools

import (
	"context"
	"sync"

	"github.com/google/uuid"

	"github.com/nextlevelbuilder/goclaw/internal/config"
	"github.com/nextlevelbuilder/goclaw/internal/sandbox"
)

// Tool execution context keys.
// These replace mutable setter fields on tool instances, making tools thread-safe
// for concurrent execution. Values are injected into context by the registry
// and read by individual tools during Execute().

type toolContextKey string

const (
	ctxChannel     toolContextKey = "tool_channel"
	ctxChannelType toolContextKey = "tool_channel_type"
	ctxChatID      toolContextKey = "tool_chat_id"
	ctxPeerKind    toolContextKey = "tool_peer_kind"
	ctxLocalKey    toolContextKey = "tool_local_key" // composite key with topic/thread suffix for routing
	ctxSandboxKey  toolContextKey = "tool_sandbox_key"
	ctxAsyncCB     toolContextKey = "tool_async_cb"
	ctxWorkspace   toolContextKey = "tool_workspace"
	ctxAgentKey    toolContextKey = "tool_agent_key"
	ctxSessionKey  toolContextKey = "tool_session_key" // origin session key for announce routing
)

// Well-known channel names used for routing and access control.
const (
	ChannelSystem    = "system"
	ChannelDashboard = "dashboard"
	ChannelTeammate  = "teammate"
)

// MediaPathLoader resolves a media ID to a local file path.
// Used by media analysis tools (read_document, read_audio, read_video).
type MediaPathLoader interface {
	LoadPath(id string) (string, error)
}

func WithToolChannel(ctx context.Context, channel string) context.Context {
	return context.WithValue(ctx, ctxChannel, channel)
}

func ToolChannelFromCtx(ctx context.Context) string {
	v, _ := ctx.Value(ctxChannel).(string)
	return v
}

func WithToolChannelType(ctx context.Context, channelType string) context.Context {
	return context.WithValue(ctx, ctxChannelType, channelType)
}

func ToolChannelTypeFromCtx(ctx context.Context) string {
	v, _ := ctx.Value(ctxChannelType).(string)
	return v
}

func WithToolChatID(ctx context.Context, chatID string) context.Context {
	return context.WithValue(ctx, ctxChatID, chatID)
}

func ToolChatIDFromCtx(ctx context.Context) string {
	v, _ := ctx.Value(ctxChatID).(string)
	return v
}

func WithToolPeerKind(ctx context.Context, peerKind string) context.Context {
	return context.WithValue(ctx, ctxPeerKind, peerKind)
}

func ToolPeerKindFromCtx(ctx context.Context) string {
	v, _ := ctx.Value(ctxPeerKind).(string)
	return v
}

// WithToolLocalKey injects the composite local key (e.g. "-100123:topic:42") into context.
// Used by delegation/subagent to preserve topic routing info for announce-back.
func WithToolLocalKey(ctx context.Context, localKey string) context.Context {
	return context.WithValue(ctx, ctxLocalKey, localKey)
}

func ToolLocalKeyFromCtx(ctx context.Context) string {
	v, _ := ctx.Value(ctxLocalKey).(string)
	return v
}

func WithToolSandboxKey(ctx context.Context, key string) context.Context {
	return context.WithValue(ctx, ctxSandboxKey, key)
}

func ToolSandboxKeyFromCtx(ctx context.Context) string {
	v, _ := ctx.Value(ctxSandboxKey).(string)
	return v
}

func WithToolAsyncCB(ctx context.Context, cb AsyncCallback) context.Context {
	return context.WithValue(ctx, ctxAsyncCB, cb)
}

func ToolAsyncCBFromCtx(ctx context.Context) AsyncCallback {
	v, _ := ctx.Value(ctxAsyncCB).(AsyncCallback)
	return v
}

func WithToolWorkspace(ctx context.Context, ws string) context.Context {
	return context.WithValue(ctx, ctxWorkspace, ws)
}

func ToolWorkspaceFromCtx(ctx context.Context) string {
	v, _ := ctx.Value(ctxWorkspace).(string)
	return v
}

// WithToolAgentKey injects the calling agent's key into context.
// Multiple agents share a single tool registry; the agent key
// lets tools like spawn/subagent identify which agent is the parent.
func WithToolAgentKey(ctx context.Context, key string) context.Context {
	return context.WithValue(ctx, ctxAgentKey, key)
}

func ToolAgentKeyFromCtx(ctx context.Context) string {
	v, _ := ctx.Value(ctxAgentKey).(string)
	return v
}

// WithToolSessionKey injects the parent's session key so subagent announce
// can route results back to the exact same session (required for WS where
// session keys don't follow BuildScopedSessionKey format).
func WithToolSessionKey(ctx context.Context, key string) context.Context {
	return context.WithValue(ctx, ctxSessionKey, key)
}

func ToolSessionKeyFromCtx(ctx context.Context) string {
	v, _ := ctx.Value(ctxSessionKey).(string)
	return v
}

// --- Builtin tool settings (global DB overrides) ---

const ctxBuiltinToolSettings toolContextKey = "tool_builtin_settings"

// BuiltinToolSettings maps tool name → settings JSON bytes.
type BuiltinToolSettings map[string][]byte

func WithBuiltinToolSettings(ctx context.Context, settings BuiltinToolSettings) context.Context {
	return context.WithValue(ctx, ctxBuiltinToolSettings, settings)
}

func BuiltinToolSettingsFromCtx(ctx context.Context) BuiltinToolSettings {
	v, _ := ctx.Value(ctxBuiltinToolSettings).(BuiltinToolSettings)
	return v
}

// --- Per-agent restrict_to_workspace override ---

const ctxRestrictWs toolContextKey = "tool_restrict_to_workspace"

// WithRestrictToWorkspace injects a per-agent restrict_to_workspace override into context.
func WithRestrictToWorkspace(ctx context.Context, restrict bool) context.Context {
	return context.WithValue(ctx, ctxRestrictWs, restrict)
}

// RestrictFromCtx returns the per-agent restrict_to_workspace override.
func RestrictFromCtx(ctx context.Context) (bool, bool) {
	v, ok := ctx.Value(ctxRestrictWs).(bool)
	return v, ok
}

func effectiveRestrict(ctx context.Context, toolDefault bool) bool {
	if v, ok := RestrictFromCtx(ctx); ok {
		return v
	}
	return toolDefault
}

// --- Parent agent model (for subagent inheritance) ---

const ctxParentModel toolContextKey = "tool_parent_model"

// WithParentModel sets the parent agent's model in context so subagents can inherit it.
func WithParentModel(ctx context.Context, model string) context.Context {
	return context.WithValue(ctx, ctxParentModel, model)
}

// ParentModelFromCtx returns the parent agent's model from context.
func ParentModelFromCtx(ctx context.Context) string {
	v, _ := ctx.Value(ctxParentModel).(string)
	return v
}

// --- Parent agent provider (for subagent inheritance) ---

const ctxParentProvider toolContextKey = "tool_parent_provider"

// WithParentProvider sets the parent agent's provider name in context so subagents inherit it.
func WithParentProvider(ctx context.Context, providerName string) context.Context {
	return context.WithValue(ctx, ctxParentProvider, providerName)
}

// ParentProviderFromCtx returns the parent agent's provider name from context.
func ParentProviderFromCtx(ctx context.Context) string {
	v, _ := ctx.Value(ctxParentProvider).(string)
	return v
}

// --- Per-agent subagent config override ---

const ctxSubagentCfg toolContextKey = "tool_subagent_config"

func WithSubagentConfig(ctx context.Context, cfg *config.SubagentsConfig) context.Context {
	return context.WithValue(ctx, ctxSubagentCfg, cfg)
}

func SubagentConfigFromCtx(ctx context.Context) *config.SubagentsConfig {
	v, _ := ctx.Value(ctxSubagentCfg).(*config.SubagentsConfig)
	return v
}

// --- Per-agent memory config override ---

const ctxMemoryCfg toolContextKey = "tool_memory_config"

func WithMemoryConfig(ctx context.Context, cfg *config.MemoryConfig) context.Context {
	return context.WithValue(ctx, ctxMemoryCfg, cfg)
}

func MemoryConfigFromCtx(ctx context.Context) *config.MemoryConfig {
	v, _ := ctx.Value(ctxMemoryCfg).(*config.MemoryConfig)
	return v
}

// --- Team ID propagation (task dispatch → workspace tools) ---

const ctxTeamID toolContextKey = "tool_team_id"

// WithToolTeamID injects the dispatching team's ID into context so team
// tools (team_tasks, team_message) and the WorkspaceInterceptor resolve
// the correct team when the agent belongs to multiple teams.
func WithToolTeamID(ctx context.Context, teamID string) context.Context {
	return context.WithValue(ctx, ctxTeamID, teamID)
}

// ToolTeamIDFromCtx returns the dispatching team's ID from context.
func ToolTeamIDFromCtx(ctx context.Context) string {
	v, _ := ctx.Value(ctxTeamID).(string)
	return v
}

// --- Team workspace path (accessible but not default) ---

const ctxTeamWorkspace toolContextKey = "tool_team_workspace"

// WithToolTeamWorkspace stores the team shared workspace directory path.
// File tools allow access to this path even when restrict_to_workspace is true.
func WithToolTeamWorkspace(ctx context.Context, dir string) context.Context {
	return context.WithValue(ctx, ctxTeamWorkspace, dir)
}

// ToolTeamWorkspaceFromCtx returns the team shared workspace directory path.
func ToolTeamWorkspaceFromCtx(ctx context.Context) string {
	v, _ := ctx.Value(ctxTeamWorkspace).(string)
	return v
}

// --- Team task ID propagation (delegation origin → workspace tools) ---

const ctxTeamTaskID toolContextKey = "tool_team_task_id"

// WithTeamTaskID injects the delegation's team task ID into context
// so workspace tools can auto-link files to the active task.
func WithTeamTaskID(ctx context.Context, taskID string) context.Context {
	return context.WithValue(ctx, ctxTeamTaskID, taskID)
}

// TeamTaskIDFromCtx returns the delegation's team task ID from context.
func TeamTaskIDFromCtx(ctx context.Context) string {
	v, _ := ctx.Value(ctxTeamTaskID).(string)
	return v
}

// --- Workspace scope propagation (delegation origin) ---

const (
	ctxWsChannel toolContextKey = "tool_workspace_channel"
	ctxWsChatID  toolContextKey = "tool_workspace_chat_id"
)

func WithWorkspaceChannel(ctx context.Context, channel string) context.Context {
	return context.WithValue(ctx, ctxWsChannel, channel)
}

func WorkspaceChannelFromCtx(ctx context.Context) string {
	v, _ := ctx.Value(ctxWsChannel).(string)
	return v
}

func WithWorkspaceChatID(ctx context.Context, chatID string) context.Context {
	return context.WithValue(ctx, ctxWsChatID, chatID)
}

func WorkspaceChatIDFromCtx(ctx context.Context) string {
	v, _ := ctx.Value(ctxWsChatID).(string)
	return v
}

// --- Pending team task dispatch (post-turn processing) ---

const ctxPendingDispatch toolContextKey = "tool_pending_team_dispatch"

// PendingTeamDispatch tracks team tasks created during an agent turn.
// After the turn ends, the consumer drains and dispatches them.
// Thread-safe: tools may execute in parallel goroutines.
type PendingTeamDispatch struct {
	mu       sync.Mutex
	tasks    map[uuid.UUID][]uuid.UUID // teamID → []taskID
	listed   bool                      // true after list called in this turn
	teamLock *sync.Mutex               // acquired on list, released before post-turn dispatch
}

func NewPendingTeamDispatch() *PendingTeamDispatch {
	return &PendingTeamDispatch{tasks: make(map[uuid.UUID][]uuid.UUID)}
}

// Add records a task created during this turn.
func (p *PendingTeamDispatch) Add(teamID, taskID uuid.UUID) {
	p.mu.Lock()
	p.tasks[teamID] = append(p.tasks[teamID], taskID)
	p.mu.Unlock()
}

// Drain returns all tracked tasks and resets the container.
func (p *PendingTeamDispatch) Drain() map[uuid.UUID][]uuid.UUID {
	p.mu.Lock()
	defer p.mu.Unlock()
	out := p.tasks
	p.tasks = make(map[uuid.UUID][]uuid.UUID)
	return out
}

// MarkListed records that list was called in this turn.
func (p *PendingTeamDispatch) MarkListed() {
	p.mu.Lock()
	p.listed = true
	p.mu.Unlock()
}

// HasListed reports whether list was called in this turn.
func (p *PendingTeamDispatch) HasListed() bool {
	p.mu.Lock()
	defer p.mu.Unlock()
	return p.listed
}

// SetTeamLock stores the acquired team create lock so it can be released post-turn.
func (p *PendingTeamDispatch) SetTeamLock(m *sync.Mutex) {
	p.mu.Lock()
	p.teamLock = m
	p.mu.Unlock()
}

// ReleaseTeamLock releases the held team create lock, if any.
func (p *PendingTeamDispatch) ReleaseTeamLock() {
	p.mu.Lock()
	if p.teamLock != nil {
		p.teamLock.Unlock()
		p.teamLock = nil
	}
	p.mu.Unlock()
}

func WithPendingTeamDispatch(ctx context.Context, ptd *PendingTeamDispatch) context.Context {
	return context.WithValue(ctx, ctxPendingDispatch, ptd)
}

func PendingTeamDispatchFromCtx(ctx context.Context) *PendingTeamDispatch {
	v, _ := ctx.Value(ctxPendingDispatch).(*PendingTeamDispatch)
	return v
}

// --- Run media file paths (for team workspace auto-collect) ---

const ctxRunMediaPaths toolContextKey = "tool_run_media_paths"

// WithRunMediaPaths stores the absolute file paths of media files received
// in the current run. Used by team_tasks to auto-copy files to team workspace.
func WithRunMediaPaths(ctx context.Context, paths []string) context.Context {
	return context.WithValue(ctx, ctxRunMediaPaths, paths)
}

// RunMediaPathsFromCtx returns media file paths from the current run.
func RunMediaPathsFromCtx(ctx context.Context) []string {
	v, _ := ctx.Value(ctxRunMediaPaths).([]string)
	return v
}

const ctxRunMediaNames toolContextKey = "tool_run_media_names"

// WithRunMediaNames stores the mapping from media file path to original filename.
func WithRunMediaNames(ctx context.Context, names map[string]string) context.Context {
	return context.WithValue(ctx, ctxRunMediaNames, names)
}

// RunMediaNamesFromCtx returns the media path → original filename mapping.
func RunMediaNamesFromCtx(ctx context.Context) map[string]string {
	v, _ := ctx.Value(ctxRunMediaNames).(map[string]string)
	return v
}

// --- Per-agent sandbox config override ---

const ctxSandboxCfg toolContextKey = "tool_sandbox_config"

func WithSandboxConfig(ctx context.Context, cfg *sandbox.Config) context.Context {
	return context.WithValue(ctx, ctxSandboxCfg, cfg)
}

func SandboxConfigFromCtx(ctx context.Context) *sandbox.Config {
	v, _ := ctx.Value(ctxSandboxCfg).(*sandbox.Config)
	return v
}