Env.hs

{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE RankNTypes #-}

-- | Execution-environment operations for built-in work tools.
--
-- Filesystem and process-backed tools should depend on this record instead of
-- calling the host operating system directly. A future sandbox can provide a
-- different 'ToolEnv' value while the tool definitions stay unchanged.
--
-- Security posture: 'localToolEnv' is deliberately non-hermetic. It has full
-- access to the host filesystem paths and inherited environment visible to the
-- current process, and it can execute arbitrary shell commands with the current
-- user's privileges. Its filesystem and process IO is performed inside the
-- @(LLM, Error ShikumiError)@ row via 'unsafeEff_'. Supply a different 'ToolEnv'
-- to confine tools; this module is the sandboxing seam, not the sandbox.
module Shikumi.Tool.Env
  ( Path,
    EnvRow,
    ExecRequest (..),
    ExecResult (..),
    FileStat (..),
    DirEntry (..),
    ToolEnv (..),
    localToolEnv,
  )
where

import Control.Exception (IOException, try)
import Control.Lens ((^.))
import Data.Aeson (ToJSON)
import Data.ByteString (ByteString)
import Data.ByteString qualified as BS
import Data.Generics.Labels ()
import Data.Text (Text)
import Data.Text qualified as T
import Effectful (Eff, (:>))
import Effectful.Dispatch.Static (unsafeEff_)
import Effectful.Error.Static (Error, throwError)
import GHC.Generics (Generic)
import Shikumi.Error (ShikumiError (..))
import Shikumi.LLM (LLM)
import System.Directory qualified as Dir
import System.Exit (ExitCode (..))
import System.Process qualified as Proc
import System.Timeout qualified as Timeout

type Path = Text

type EnvRow es = (LLM :> es, Error ShikumiError :> es)

data ExecRequest = ExecRequest
  { command :: !Text,
    cwd :: !(Maybe Path),
    stdin :: !(Maybe Text),
    timeoutMs :: !(Maybe Int)
  }
  deriving stock (Generic, Show, Eq)
  deriving anyclass (ToJSON)

data ExecResult = ExecResult
  { exitCode :: !Int,
    stdout :: !Text,
    stderr :: !Text
  }
  deriving stock (Generic, Show, Eq)
  deriving anyclass (ToJSON)

data FileStat = FileStat
  { isFile :: !Bool,
    isDir :: !Bool,
    size :: !Integer
  }
  deriving stock (Generic, Show, Eq)
  deriving anyclass (ToJSON)

data DirEntry = DirEntry
  { name :: !Text,
    isDir :: !Bool,
    isSymlink :: !Bool
  }
  deriving stock (Generic, Show, Eq)
  deriving anyclass (ToJSON)

data ToolEnv = ToolEnv
  { envExec :: forall es. (EnvRow es) => ExecRequest -> Eff es ExecResult,
    envReadFile :: forall es. (EnvRow es) => Path -> Eff es ByteString,
    envWriteFile :: forall es. (EnvRow es) => Path -> ByteString -> Eff es (),
    envStat :: forall es. (EnvRow es) => Path -> Eff es (Maybe FileStat),
    envReaddir :: forall es. (EnvRow es) => Path -> Eff es [DirEntry],
    envExists :: forall es. (EnvRow es) => Path -> Eff es Bool,
    envMkdir :: forall es. (EnvRow es) => Path -> Eff es (),
    envRm :: forall es. (EnvRow es) => Path -> Eff es (),
    envCwd :: forall es. (EnvRow es) => Eff es Path
  }

localToolEnv :: ToolEnv
localToolEnv =
  ToolEnv
    { envExec = localExec,
      envReadFile = \path -> toolIO "readFile" (BS.readFile (T.unpack path)),
      envWriteFile = \path bytes -> toolIO "writeFile" (BS.writeFile (T.unpack path) bytes),
      envStat = localStat,
      envReaddir = localReaddir,
      envExists = \path -> toolIO "exists" (Dir.doesPathExist (T.unpack path)),
      envMkdir = \path -> toolIO "mkdir" (Dir.createDirectoryIfMissing True (T.unpack path)),
      envRm = localRm,
      envCwd = T.pack <$> toolIO "cwd" Dir.getCurrentDirectory
    }

localExec :: (EnvRow es) => ExecRequest -> Eff es ExecResult
localExec req = do
  let effectiveTimeoutMs = effectiveTimeout (req ^. #timeoutMs)
  result <-
    toolIO "exec" $
      Timeout.timeout (effectiveTimeoutMs * 1000) $
        Proc.readCreateProcessWithExitCode process stdinText
  case result of
    Nothing ->
      throwError $
        Timeout
          ("tool env: exec timed out after " <> T.pack (show effectiveTimeoutMs) <> "ms")
    Just (exit, out, err) ->
      pure
        ExecResult
          { exitCode = exitCodeInt exit,
            stdout = T.pack out,
            stderr = T.pack err
          }
  where
    process =
      (Proc.proc "bash" ["-c", T.unpack (req ^. #command)])
        { Proc.cwd = T.unpack <$> (req ^. #cwd)
        }
    stdinText = maybe "" T.unpack (req ^. #stdin)

localStat :: (EnvRow es) => Path -> Eff es (Maybe FileStat)
localStat path = toolIO "stat" $ do
  let fp = T.unpack path
  file <- Dir.doesFileExist fp
  dir <- Dir.doesDirectoryExist fp
  if file
    then do
      size <- Dir.getFileSize fp
      pure (Just FileStat {isFile = True, isDir = False, size})
    else
      if dir
        then pure (Just FileStat {isFile = False, isDir = True, size = 0})
        else pure Nothing

localReaddir :: (EnvRow es) => Path -> Eff es [DirEntry]
localReaddir path = toolIO "readdir" $ do
  let dir = T.unpack path
  names <- Dir.listDirectory dir
  traverse
    ( \entry -> do
        let child = dir <> "/" <> entry
        isSymlink <- Dir.pathIsSymbolicLink child
        isDir <- Dir.doesDirectoryExist child
        pure DirEntry {name = T.pack entry, isDir, isSymlink}
    )
    names

localRm :: (EnvRow es) => Path -> Eff es ()
localRm path = toolIO "rm" $ do
  let fp = T.unpack path
  file <- Dir.doesFileExist fp
  dir <- Dir.doesDirectoryExist fp
  if file
    then Dir.removeFile fp
    else
      if dir
        then Dir.removeDirectory fp
        else pure ()

toolIO :: (EnvRow es) => Text -> IO a -> Eff es a
toolIO label action = do
  result <- unsafeEff_ (try action)
  case result of
    Right value -> pure value
    Left (err :: IOException) ->
      throwError (ProviderFailure ("tool env: " <> label <> ": " <> T.pack (show err)))

-- | Clamp a model-supplied timeout into [1 ms, 'maxExecTimeoutMs']. Unclamped,
-- a non-positive value disables the timeout entirely and a huge value can
-- overflow the microsecond multiplication in 'localExec'.
effectiveTimeout :: Maybe Int -> Int
effectiveTimeout = max 1 . min maxExecTimeoutMs . maybe defaultExecTimeoutMs id

defaultExecTimeoutMs :: Int
defaultExecTimeoutMs = 60000

maxExecTimeoutMs :: Int
maxExecTimeoutMs = 600000

exitCodeInt :: ExitCode -> Int
exitCodeInt ExitSuccess = 0
exitCodeInt (ExitFailure n) = n