diff --git a/package.json b/package.json
index 7eb34cc..66fe1a9 100644
--- a/package.json
+++ b/package.json
@@ -3,7 +3,7 @@
   "publishConfig": {
     "access": "public"
   },
-  "version": "1.3.8",
+  "version": "1.3.9",
   "description": "Generic bash tool for AI agents, compatible with AI SDK",
   "type": "module",
   "main": "dist/index.js",
diff --git a/src/files/loader.ts b/src/files/loader.ts
index 791c19d..2ae07ed 100644
--- a/src/files/loader.ts
+++ b/src/files/loader.ts
@@ -63,6 +63,7 @@ export async function* streamFiles(
 export async function getFilePaths(
   options: LoadFilesOptions,
 ): Promise<string[]> {
+  "use step";
   const paths: string[] = [];
 
   if (options.uploadDirectory) {
diff --git a/src/sandbox/just-bash.ts b/src/sandbox/just-bash.ts
index 40f2d12..1927139 100644
--- a/src/sandbox/just-bash.ts
+++ b/src/sandbox/just-bash.ts
@@ -14,6 +14,8 @@ export interface JustBashLike {
     readFile: (path: string) => Promise<string>;
     writeFile: (path: string, content: string) => Promise<void>;
   };
+  limits: unknown;
+  state: unknown;
 }
 
 /**
diff --git a/src/tool.ts b/src/tool.ts
index 655e485..97a54a4 100644
--- a/src/tool.ts
+++ b/src/tool.ts
@@ -3,6 +3,7 @@ import { getFilePaths, streamFiles } from "./files/loader.js";
 import {
   createJustBashSandbox,
   isJustBash,
+  JustBashLike,
   wrapJustBash,
 } from "./sandbox/just-bash.js";
 import { isVercelSandbox, wrapVercelSandbox } from "./sandbox/vercel.js";
@@ -10,7 +11,13 @@ import { createBashExecuteTool } from "./tools/bash.js";
 import { createReadFileTool } from "./tools/read-file.js";
 import { createWriteFileTool } from "./tools/write-file.js";
 import { createToolPrompt } from "./tools-prompt.js";
-import type { BashToolkit, CreateBashToolOptions, Sandbox } from "./types.js";
+import type {
+  BashToolkit,
+  CommandResult,
+  CreateBashToolOptions,
+  Sandbox,
+} from "./types.js";
+import { ToolExecuteFunction } from "ai";
 
 const DEFAULT_DESTINATION = "/workspace";
 const VERCEL_SANDBOX_DESTINATION = "/vercel/sandbox/workspace";
@@ -44,8 +51,28 @@ const DEFAULT_MAX_FILES = 1000;
  * ```
  */
 export async function createBashTool(
-  options: CreateBashToolOptions = {},
+  options: CreateBashToolOptions = {}
 ): Promise<BashToolkit> {
+  const tools = await createBashToolStep(options);
+  const bashStep = tools.bash.execute!;
+  tools.bash.execute = async (input, toolOptions) => {
+    const result = (await bashStep(input, toolOptions)) as CommandResult;
+    // For just-bash we need to copy the FS from the returned sandbox
+    if (isJustBash(result.sandbox) && isJustBash(options.sandbox)) {
+      const sandbox = options.sandbox as JustBashLike;
+      sandbox.fs = result.sandbox.fs;
+      sandbox.limits = result.sandbox.limits;
+      sandbox.state = result.sandbox.state;
+    }
+    return result;
+  };
+  return tools;
+}
+
+export async function createBashToolStep(
+  options: CreateBashToolOptions = {}
+): Promise<BashToolkit> {
+  "use step";
   // Determine default destination based on sandbox type
   const defaultDestination =
     options.sandbox && isVercelSandbox(options.sandbox)
@@ -86,7 +113,7 @@ export async function createBashTool(
       throw new Error(
         `Too many files to upload: ${fileList.length} files exceeds the limit of ${maxFiles}. ` +
           `Either increase maxFiles, use a more restrictive include pattern in uploadDirectory, ` +
-          `or write files to the sandbox yourself before calling createBashTool.`,
+          `or write files to the sandbox yourself before calling createBashTool.`
       );
     }
 
@@ -135,7 +162,7 @@ export async function createBashTool(
       if (maxFiles > 0 && fileList.length > maxFiles) {
         throw new Error(
           `Too many files: ${fileList.length} files exceeds the limit of ${maxFiles}. ` +
-            `Either increase maxFiles or use a more restrictive include pattern in uploadDirectory.`,
+            `Either increase maxFiles or use a more restrictive include pattern in uploadDirectory.`
         );
       }
 
@@ -166,7 +193,7 @@ export async function createBashTool(
         throw new Error(
           `Too many files to load: ${fileList.length} files exceeds the limit of ${maxFiles}. ` +
             `Either increase maxFiles, use a more restrictive include pattern in uploadDirectory, ` +
-            `or provide your own sandbox with files already written.`,
+            `or provide your own sandbox with files already written.`
         );
       }
 
diff --git a/src/tools/bash.ts b/src/tools/bash.ts
index 813cb41..1e5cda4 100644
--- a/src/tools/bash.ts
+++ b/src/tools/bash.ts
@@ -114,6 +114,7 @@ export function createBashExecuteTool(options: CreateBashToolOptions) {
     description: generateDescription(options),
     inputSchema: bashSchema,
     execute: async ({ command: originalCommand }) => {
+      "use step";
       // Allow modification of command before execution
       let command = originalCommand;
       if (onBeforeBashCall) {
@@ -134,6 +135,7 @@ export function createBashExecuteTool(options: CreateBashToolOptions) {
         ...result,
         stdout: truncateOutput(result.stdout, maxOutputLength, "stdout"),
         stderr: truncateOutput(result.stderr, maxOutputLength, "stderr"),
+        sandbox,
       };
 
       // Allow modification of result after execution
diff --git a/src/tools/read-file.ts b/src/tools/read-file.ts
index 63f47b9..144955a 100644
--- a/src/tools/read-file.ts
+++ b/src/tools/read-file.ts
@@ -20,6 +20,7 @@ export function createReadFileTool(options: CreateReadFileToolOptions) {
     description: "Read the contents of a file from the sandbox.",
     inputSchema: readFileSchema,
     execute: async ({ path }) => {
+      "use step";
       const resolvedPath = nodePath.posix.resolve(cwd, path);
       const content = await sandbox.readFile(resolvedPath);
       return { content };
diff --git a/src/tools/write-file.ts b/src/tools/write-file.ts
index 6a31363..69be95c 100644
--- a/src/tools/write-file.ts
+++ b/src/tools/write-file.ts
@@ -22,6 +22,7 @@ export function createWriteFileTool(options: CreateWriteFileToolOptions) {
       "Write content to a file in the sandbox. Creates parent directories if needed.",
     inputSchema: writeFileSchema,
     execute: async ({ path, content }) => {
+      "use step";
       const resolvedPath = nodePath.posix.resolve(cwd, path);
       await sandbox.writeFiles([{ path: resolvedPath, content }]);
       return { success: true };
diff --git a/src/types.ts b/src/types.ts
index 37a7824..2c2ad4f 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -5,6 +5,7 @@ export interface CommandResult {
   stdout: string;
   stderr: string;
   exitCode: number;
+  sandbox: Sandbox;
 }
 
 export interface Sandbox {