diff --git a/packages/runtime/src/constants.ts b/packages/runtime/src/constants.ts
index 3c12b5a88..36143621f 100644
--- a/packages/runtime/src/constants.ts
+++ b/packages/runtime/src/constants.ts
@@ -53,11 +53,6 @@ export enum PrismaErrorCode {
     DEPEND_ON_RECORD_NOT_FOUND = 'P2025',
 }
 
-/**
- * Field name for storing in-transaction flag
- */
-export const PRISMA_TX_FLAG = '$__zenstack_tx';
-
 /**
  * Field name for getting current enhancer
  */
diff --git a/packages/runtime/src/enhancements/policy/handler.ts b/packages/runtime/src/enhancements/policy/handler.ts
index f002002d2..9268c24b3 100644
--- a/packages/runtime/src/enhancements/policy/handler.ts
+++ b/packages/runtime/src/enhancements/policy/handler.ts
@@ -4,7 +4,7 @@ import { lowerCaseFirst } from 'lower-case-first';
 import invariant from 'tiny-invariant';
 import { upperCaseFirst } from 'upper-case-first';
 import { fromZodError } from 'zod-validation-error';
-import { CrudFailureReason, PRISMA_TX_FLAG } from '../../constants';
+import { CrudFailureReason } from '../../constants';
 import {
     ModelDataVisitor,
     NestedWriteVisitor,
@@ -1232,11 +1232,11 @@ export class PolicyProxyHandler<DbClient extends DbClientContract> implements Pr
     }
 
     private transaction(action: (tx: Record<string, DbOperations>) => Promise<any>) {
-        if (this.prisma[PRISMA_TX_FLAG]) {
+        if (this.prisma['$transaction']) {
+            return this.prisma.$transaction((tx) => action(tx), { maxWait: 100000, timeout: 100000 });
+        } else {
             // already in transaction, don't nest
             return action(this.prisma);
-        } else {
-            return this.prisma.$transaction((tx) => action(tx), { maxWait: 100000, timeout: 100000 });
         }
     }
 
diff --git a/packages/runtime/src/enhancements/proxy.ts b/packages/runtime/src/enhancements/proxy.ts
index 358bff153..c735d595a 100644
--- a/packages/runtime/src/enhancements/proxy.ts
+++ b/packages/runtime/src/enhancements/proxy.ts
@@ -1,6 +1,6 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
 
-import { PRISMA_PROXY_ENHANCER, PRISMA_TX_FLAG } from '../constants';
+import { PRISMA_PROXY_ENHANCER } from '../constants';
 import type { ModelMeta } from '../cross';
 import type { DbClientContract } from '../types';
 import { createDeferredPromise } from './policy/promise';
@@ -183,6 +183,7 @@ export function makeProxy<T extends PrismaProxyHandler>(
     errorTransformer?: ErrorTransformer
 ) {
     const models = Object.keys(modelMeta.fields).map((k) => k.toLowerCase());
+
     const proxy = new Proxy(prisma, {
         get: (target: any, prop: string | symbol, receiver: any) => {
             // enhancer metadata
@@ -191,7 +192,7 @@ export function makeProxy<T extends PrismaProxyHandler>(
             }
 
             if (prop === 'toString') {
-                return () => `$zenstack_${name}[${target.toString()}]`;
+                return () => `$zenstack_prisma_${prisma._clientVersion}`;
             }
 
             if (prop === '$transaction') {
@@ -213,8 +214,10 @@ export function makeProxy<T extends PrismaProxyHandler>(
 
                         const txFunc = input;
                         return $transaction.bind(target)((tx: any) => {
+                            // create a proxy for the transaction function
                             const txProxy = makeProxy(tx, modelMeta, makeHandler, name + '$tx');
-                            txProxy[PRISMA_TX_FLAG] = true;
+
+                            // call the transaction function with the proxy
                             return txFunc(txProxy);
                         }, ...rest);
                     };
diff --git a/tests/integration/tests/enhancements/with-policy/postgres.test.ts b/tests/integration/tests/enhancements/with-policy/postgres.test.ts
index caed6a5ce..7653b89f7 100644
--- a/tests/integration/tests/enhancements/with-policy/postgres.test.ts
+++ b/tests/integration/tests/enhancements/with-policy/postgres.test.ts
@@ -16,14 +16,14 @@ describe('With Policy: with postgres', () => {
 
     beforeEach(async () => {
         dbUrl = await createPostgresDb(DB_NAME);
-        const { prisma: _prisma, withPolicy } = await loadSchemaFromFile(
+        const { prisma: _prisma, enhance } = await loadSchemaFromFile(
             path.join(__dirname, '../../schema/todo-pg.zmodel'),
             {
                 provider: 'postgresql',
                 dbUrl,
             }
         );
-        getDb = withPolicy;
+        getDb = enhance;
         prisma = _prisma;
     });