diff --git a/client/packages/cli/__tests__/__snapshots__/updateSchemaFile.test.ts.snap b/client/packages/cli/__tests__/__snapshots__/updateSchemaFile.test.ts.snap
new file mode 100644
index 0000000000..e935e58dd7
--- /dev/null
+++ b/client/packages/cli/__tests__/__snapshots__/updateSchemaFile.test.ts.snap
@@ -0,0 +1,248 @@
+// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
+
+exports[`adds a link when links object is empty 1`] = `
+"
+import { i } from '@instantdb/core';
+
+const _schema = i.schema({
+  entities: {
+    todos: i.entity({
+      title: i.string(),
+    }),
+    users: i.entity({
+      email: i.string(),
+    }),
+  },
+  links: {
+    todoOwner: {
+      forward: {
+        on: 'todos',
+        has: 'one',
+        label: 'owner',
+      },
+      reverse: {
+        on: 'users',
+        has: 'many',
+        label: 'todos',
+      },
+    },
+  },
+  rooms: {},
+});
+
+export default _schema;
+"
+`;
+
+exports[`drops constraints removed by server 1`] = `
+"
+import { i } from '@instantdb/core';
+
+const _schema = i.schema({
+  entities: {
+    todos: i.entity({
+      title: i.string(),
+      done: i.boolean().optional(),
+    }),
+  },
+  links: {},
+  rooms: {},
+});
+
+export default _schema;
+"
+`;
+
+exports[`handles quoted keys for entities, attrs, and links 1`] = `
+"
+import { i } from '@instantdb/core';
+
+const _schema = i.schema({
+  entities: {
+    todos: i.entity({
+      title: i.string(),
+    }),
+    users: i.entity({
+      email: i.string(),
+    }),
+    'user-profiles': i.entity({
+      'display-name': i.string(),
+      'avatar-url': i.string(),
+    }),
+  },
+  links: {
+    'todo-owner': {
+      forward: {
+        on: 'todos',
+        has: 'one',
+        label: 'owner',
+      },
+      reverse: {
+        on: 'users',
+        has: 'many',
+        label: 'todos',
+      },
+    },
+  },
+  rooms: {},
+});
+
+export default _schema;
+"
+`;
+
+exports[`inserts attrs into multi-line entities with indentation 1`] = `
+"
+import { i } from '@instantdb/core';
+
+const _schema = i.schema({
+  entities: {
+    todos: i.entity({
+      title: i.string(),
+      done: i.boolean().optional(),
+      priority: i.number(),
+    }),
+  },
+  links: {},
+  rooms: {},
+});
+
+export default _schema;
+"
+`;
+
+exports[`preserves type params across chained calls 1`] = `
+"
+import { i } from '@instantdb/core';
+import { Label } from './types';
+
+const _schema = i.schema({
+  entities: {
+    todos: i.entity({
+      title: i.string(),
+      status: i.string<'todo' | 'done'>().unique().indexed(),
+      labels: i.json<Label[]>(),
+      metadata: i.json(),
+    }),
+    users: i.entity({
+      email: i.string().unique(),
+    }),
+  },
+  links: {},
+  rooms: {},
+});
+
+export default _schema;
+"
+`;
+
+exports[`removes a link with surrounding comments and commas 1`] = `
+"
+import { i } from '@instantdb/core';
+
+const _schema = i.schema({
+  entities: {
+    todos: i.entity({
+      title: i.string(),
+    }),
+    users: i.entity({
+      email: i.string(),
+    }),
+    projects: i.entity({
+      name: i.string(),
+    }),
+  },
+  links: {
+    // owner link
+    /* project link */
+    projectTodos: {
+      forward: { on: 'projects', has: 'many', label: 'todos' },
+      reverse: { on: 'todos', has: 'one', label: 'project' },
+    },
+  },
+  rooms: {},
+});
+
+export default _schema;
+"
+`;
+
+exports[`removes the last link cleanly 1`] = `
+"
+import { i } from '@instantdb/core';
+
+const _schema = i.schema({
+  entities: {
+    todos: i.entity({
+      title: i.string(),
+    }),
+    users: i.entity({
+      email: i.string(),
+    }),
+  },
+  links: {},
+  rooms: {},
+});
+
+export default _schema;
+"
+`;
+
+exports[`updates link details 1`] = `
+"
+import { i } from '@instantdb/core';
+
+const _schema = i.schema({
+  entities: {
+    todos: i.entity({
+      title: i.string(),
+    }),
+    users: i.entity({
+      email: i.string(),
+    }),
+  },
+  links: {
+    todoOwner: {
+      forward: {
+        on: 'todos',
+        has: 'one',
+        label: 'owner',
+        required: true,
+        onDelete: 'cascade',
+      },
+      reverse: {
+        on: 'users',
+        has: 'many',
+        label: 'todos',
+        onDelete: 'cascade',
+      },
+    },
+  },
+  rooms: {},
+});
+
+export default _schema;
+"
+`;
+
+exports[`updates single-line entity in place 1`] = `
+"
+import { i } from '@instantdb/core';
+
+const _schema = i.schema({
+  entities: {
+    projects: i.entity({
+      name: i.string(),
+      status: i.string(),
+    }),
+    todos: i.entity({
+      title: i.string(),
+    }),
+  },
+  links: {},
+  rooms: {},
+});
+
+export default _schema;
+"
+`;
diff --git a/client/packages/cli/__tests__/mergeSchema.test.ts b/client/packages/cli/__tests__/mergeSchema.test.ts
deleted file mode 100644
index 1f4ed8989c..0000000000
--- a/client/packages/cli/__tests__/mergeSchema.test.ts
+++ /dev/null
@@ -1,197 +0,0 @@
-import { describe, test, expect } from 'vitest';
-import { mergeSchema } from '../src/util/mergeSchema';
-
-test('preserves type annotations', () => {
-  const oldFile = `
-import { i } from '@instantdb/core';
-import { Label } from './types';
-
-const _schema = i.schema({
-  entities: {
-    $users: i.entity({
-      email: i.string().unique().indexed(),
-    }),
-    todos: i.entity({
-      title: i.string(),
-      status: i.string<'todo' | 'in_progress' | 'done'>(),
-      priority: i.number<1 | 2 | 3>(),
-      labels: i.json<Label[]>().optional(),
-    }),
-    projects: i.entity({
-      name: i.string(),
-    }),
-  },
-  links: {
-    todoProject: {
-      forward: { on: 'todos', has: 'one', label: 'project' },
-      reverse: { on: 'projects', has: 'many', label: 'todos' },
-    },
-    projectOwner: {
-      forward: { on: 'projects', has: 'one', label: 'owner' },
-      reverse: { on: '$users', has: 'many', label: 'projects' },
-    },
-  },
-  rooms: {
-    projectRoom: {
-      presence: i.entity({
-        cursor: i.json<{ x: number; y: number }>(),
-      }),
-    },
-  },
-});
-
-export default _schema;
-`;
-
-  const newFile = `
-import { i } from '@instantdb/core';
-
-const _schema = i.schema({
-  entities: {
-    $users: i.entity({
-      email: i.string().unique().indexed(),
-    }),
-    todos: i.entity({
-      title: i.string(),
-      status: i.string(),
-      priority: i.number(),
-      labels: i.json().optional(),
-    }),
-    projects: i.entity({
-      name: i.string(),
-    }),
-  },
-  links: {
-    todoProject: {
-      forward: { on: 'todos', has: 'one', label: 'project' },
-      reverse: { on: 'projects', has: 'many', label: 'todos' },
-    },
-    projectOwner: {
-      forward: { on: 'projects', has: 'one', label: 'owner' },
-      reverse: { on: '$users', has: 'many', label: 'projects' },
-    },
-  },
-  rooms: {
-    projectRoom: {
-      presence: i.entity({
-        cursor: i.json(),
-      }),
-    },
-  },
-});
-
-export default _schema;
-`;
-
-  const result = mergeSchema(oldFile, newFile);
-
-  // Type annotations preserved
-  expect(result).toContain("i.string<'todo' | 'in_progress' | 'done'>()");
-  expect(result).toContain('i.number<1 | 2 | 3>()');
-  expect(result).toContain('i.json<Label[]>()');
-  expect(result).toContain('i.json<{ x: number; y: number }>()');
-
-  // Import preserved
-  expect(result).toContain("import { Label } from './types';");
-});
-
-test('preserves different import styles', () => {
-  const oldFile = `
-import { i } from '@instantdb/core';
-import { Label } from './types';
-import { Tag as MyTag } from './types';
-import Priority from './Priority';
-import * as Models from './models';
-import type { Meta } from './meta';
-
-const _schema = i.schema({
-  entities: {
-    a: i.entity({ f: i.json<Label>() }),
-    b: i.entity({ f: i.json<MyTag>() }),
-    c: i.entity({ f: i.json<Priority>() }),
-    d: i.entity({ f: i.json<Models.Status>() }),
-    e: i.entity({ f: i.json<Meta>() }),
-  },
-});
-`;
-
-  const newFile = `
-import { i } from '@instantdb/core';
-
-const _schema = i.schema({
-  entities: {
-    a: i.entity({ f: i.json() }),
-    b: i.entity({ f: i.json() }),
-    c: i.entity({ f: i.json() }),
-    d: i.entity({ f: i.json() }),
-    e: i.entity({ f: i.json() }),
-  },
-});
-`;
-
-  const result = mergeSchema(oldFile, newFile);
-
-  // All import styles preserved (named imports from same module are combined)
-  expect(result).toContain('Label');
-  expect(result).toContain('Tag as MyTag');
-  expect(result).toContain("import Priority from './Priority';");
-  expect(result).toContain("import * as Models from './models';");
-  expect(result).toContain("import type { Meta } from './meta';");
-
-  // Type annotations preserved
-  expect(result).toContain('i.json<Label>()');
-  expect(result).toContain('i.json<MyTag>()');
-  expect(result).toContain('i.json<Priority>()');
-  expect(result).toContain('i.json<Models.Status>()');
-  expect(result).toContain('i.json<Meta>()');
-});
-
-test('handles entity additions and removals', () => {
-  const oldFile = `
-import { i } from '@instantdb/core';
-import { Label } from './types';
-
-const _schema = i.schema({
-  entities: {
-    todos: i.entity({
-      labels: i.json<Label[]>(),
-    }),
-    oldEntity: i.entity({
-      data: i.json<{ removed: true }>(),
-    }),
-  },
-});
-`;
-
-  const newFile = `
-import { i } from '@instantdb/core';
-import { Label } from './types';
-
-const _schema = i.schema({
-  entities: {
-    todos: i.entity({
-      labels: i.json(),
-    }),
-    newEntity: i.entity({
-      data: i.json(),
-    }),
-  },
-});
-`;
-
-  const result = mergeSchema(oldFile, newFile);
-
-  // Existing entity type preserved
-  expect(result).toContain('i.json<Label[]>()');
-
-  // Removed entity is gone
-  expect(result).not.toContain('oldEntity');
-
-  // New entity has no type annotation
-  expect(result).toContain('newEntity: i.entity({');
-  expect(result).not.toContain('i.json<{ removed: true }>()');
-
-  // Import not duplicated (was already in newFile)
-  const importMatches = result.match(/import { Label } from '\.\/types';/g);
-  expect(importMatches?.length).toBe(1);
-});
diff --git a/client/packages/cli/__tests__/updateSchemaFile.test.ts b/client/packages/cli/__tests__/updateSchemaFile.test.ts
new file mode 100644
index 0000000000..29b77f2d3a
--- /dev/null
+++ b/client/packages/cli/__tests__/updateSchemaFile.test.ts
@@ -0,0 +1,557 @@
+import { test, expect } from 'vitest';
+import {
+  diffSchemas,
+  i,
+  schemaTypescriptFileToInstantSchema,
+} from '@instantdb/platform';
+import { updateSchemaFile } from '../src/util/updateSchemaFile';
+
+test('throws when schema call is missing', async () => {
+  const oldFile = `
+import { i } from '@instantdb/core';
+
+export const nope = 1;
+`;
+  const schema = i.schema({
+    entities: { todos: i.entity({ title: i.string() }) },
+    links: {},
+  });
+  await expect(updateSchemaFile(oldFile, schema, schema)).rejects.toThrow(
+    'Could not find i.schema',
+  );
+});
+
+test('throws when entities object is missing', async () => {
+  const oldFile = `
+import { i } from '@instantdb/core';
+
+const _schema = i.schema({
+  links: {},
+  rooms: {},
+});
+
+export default _schema;
+`;
+  const schema = i.schema({
+    entities: { todos: i.entity({ title: i.string() }) },
+    links: {},
+  });
+  await expect(updateSchemaFile(oldFile, schema, schema)).rejects.toThrow(
+    'entities object',
+  );
+});
+
+test('throws when links object is missing', async () => {
+  const oldFile = `
+import { i } from '@instantdb/core';
+
+const _schema = i.schema({
+  entities: {
+    todos: i.entity({
+      title: i.string(),
+    }),
+  },
+  rooms: {},
+});
+
+export default _schema;
+`;
+  const schema = i.schema({
+    entities: { todos: i.entity({ title: i.string() }) },
+    links: {},
+  });
+  await expect(updateSchemaFile(oldFile, schema, schema)).rejects.toThrow(
+    'links object',
+  );
+});
+
+test('preserves type params across chained calls', async () => {
+  const oldFile = `
+import { i } from '@instantdb/core';
+import { Label } from './types';
+
+const _schema = i.schema({
+  entities: {
+    todos: i.entity({
+      title: i.string(),
+      status: i.string<'todo' | 'done'>().optional().indexed(),
+      labels: i.json<Label[]>(),
+    }),
+    users: i.entity({
+      email: i.string(),
+    }),
+  },
+  links: {
+  },
+  rooms: {},
+});
+
+export default _schema;
+`;
+  const serverSchema = i.schema({
+    entities: {
+      todos: i.entity({
+        title: i.string(),
+        status: i.string().unique().indexed(),
+        labels: i.json(),
+        metadata: i.json(),
+      }),
+      users: i.entity({
+        email: i.string().unique(),
+      }),
+    },
+    links: {},
+  });
+
+  const result = await update(oldFile, serverSchema);
+
+  expect(result).toMatchSnapshot();
+});
+
+test('drops constraints removed by server', async () => {
+  const oldFile = `
+import { i } from '@instantdb/core';
+
+const _schema = i.schema({
+  entities: {
+    todos: i.entity({
+      title: i.string().unique().indexed().optional(),
+      done: i.boolean().optional(),
+    }),
+  },
+  links: {
+  },
+  rooms: {},
+});
+
+export default _schema;
+`;
+  const serverSchema = i.schema({
+    entities: {
+      todos: i.entity({
+        title: i.string(),
+        done: i.boolean().optional(),
+      }),
+    },
+    links: {},
+  });
+
+  const result = await update(oldFile, serverSchema);
+
+  expect(result).toMatchSnapshot();
+});
+
+test('updates link details', async () => {
+  const oldFile = `
+import { i } from '@instantdb/core';
+
+const _schema = i.schema({
+  entities: {
+    todos: i.entity({
+      title: i.string(),
+    }),
+    users: i.entity({
+      email: i.string(),
+    }),
+  },
+  links: {
+    todoOwner: {
+      forward: { on: 'todos', has: 'one', label: 'owner' },
+      reverse: { on: 'users', has: 'many', label: 'todos' },
+    },
+  },
+  rooms: {},
+});
+
+export default _schema;
+`;
+  const serverSchema = i.schema({
+    entities: {
+      todos: i.entity({ title: i.string() }),
+      users: i.entity({ email: i.string() }),
+    },
+    links: {
+      todoOwner: {
+        forward: {
+          on: 'todos',
+          has: 'one',
+          label: 'owner',
+          required: true,
+          onDelete: 'cascade',
+        },
+        reverse: {
+          on: 'users',
+          has: 'many',
+          label: 'todos',
+          onDelete: 'cascade',
+        },
+      },
+    },
+  });
+
+  const result = await update(oldFile, serverSchema);
+
+  expect(result).toMatchSnapshot();
+});
+
+test('removes a link with surrounding comments and commas', async () => {
+  const oldFile = `
+import { i } from '@instantdb/core';
+
+const _schema = i.schema({
+  entities: {
+    todos: i.entity({
+      title: i.string(),
+    }),
+    users: i.entity({
+      email: i.string(),
+    }),
+    projects: i.entity({
+      name: i.string(),
+    }),
+  },
+  links: {
+    // owner link
+    todoOwner: {
+      forward: { on: 'todos', has: 'one', label: 'owner' },
+      reverse: { on: 'users', has: 'many', label: 'todos' },
+    },
+    /* project link */
+    projectTodos: {
+      forward: { on: 'projects', has: 'many', label: 'todos' },
+      reverse: { on: 'todos', has: 'one', label: 'project' },
+    },
+  },
+  rooms: {},
+});
+
+export default _schema;
+`;
+  const serverSchema = i.schema({
+    entities: {
+      todos: i.entity({ title: i.string() }),
+      users: i.entity({ email: i.string() }),
+      projects: i.entity({ name: i.string() }),
+    },
+    links: {
+      projectTodos: {
+        forward: { on: 'projects', has: 'many', label: 'todos' },
+        reverse: { on: 'todos', has: 'one', label: 'project' },
+      },
+    },
+  });
+
+  const result = await update(oldFile, serverSchema);
+
+  expect(result).toMatchSnapshot();
+});
+
+test('updates single-line entity in place', async () => {
+  const oldFile = `
+import { i } from '@instantdb/core';
+
+const _schema = i.schema({
+  entities: {
+    projects: i.entity({ name: i.string() }),
+    todos: i.entity({
+      title: i.string(),
+    }),
+  },
+  links: {
+  },
+  rooms: {},
+});
+
+export default _schema;
+`;
+  const serverSchema = i.schema({
+    entities: {
+      projects: i.entity({
+        name: i.string(),
+        status: i.string(),
+      }),
+      todos: i.entity({
+        title: i.string(),
+      }),
+    },
+    links: {},
+  });
+
+  const result = await update(oldFile, serverSchema);
+
+  expect(result).toMatchSnapshot();
+});
+
+test('adds multiple attrs to a single-line entity', async () => {
+  const oldFile = `
+import { i } from '@instantdb/core';
+
+const _schema = i.schema({
+  entities: {
+    projects: i.entity({ name: i.string() }),
+  },
+  links: {
+  },
+  rooms: {},
+});
+
+export default _schema;
+`;
+  const serverSchema = i.schema({
+    entities: {
+      projects: i.entity({
+        name: i.string(),
+        status: i.string(),
+        priority: i.number(),
+      }),
+    },
+    links: {},
+  });
+
+  const result = await update(oldFile, serverSchema);
+
+  await expectNoDiff(result, serverSchema);
+});
+
+test('inserts attrs into multi-line entities with indentation', async () => {
+  const oldFile = `
+import { i } from '@instantdb/core';
+
+const _schema = i.schema({
+  entities: {
+    todos: i.entity({
+      title: i.string(),
+      done: i.boolean().optional(),
+    }),
+  },
+  links: {
+  },
+  rooms: {},
+});
+
+export default _schema;
+`;
+  const serverSchema = i.schema({
+    entities: {
+      todos: i.entity({
+        title: i.string(),
+        done: i.boolean().optional(),
+        priority: i.number(),
+      }),
+    },
+    links: {},
+  });
+
+  const result = await update(oldFile, serverSchema);
+
+  expect(result).toMatchSnapshot();
+});
+
+test('adds and updates attrs in a single-line entity', async () => {
+  const oldFile = `
+import { i } from '@instantdb/core';
+
+const _schema = i.schema({
+  entities: {
+    projects: i.entity({ name: i.string().optional(), code: i.string() }),
+  },
+  links: {
+  },
+  rooms: {},
+});
+
+export default _schema;
+`;
+  const serverSchema = i.schema({
+    entities: {
+      projects: i.entity({
+        name: i.string(),
+        code: i.string().unique(),
+        status: i.string(),
+        priority: i.number(),
+      }),
+    },
+    links: {},
+  });
+
+  const result = await update(oldFile, serverSchema);
+
+  await expectNoDiff(result, serverSchema);
+});
+
+test('handles quoted keys for entities, attrs, and links', async () => {
+  const oldFile = `
+import { i } from '@instantdb/core';
+
+const _schema = i.schema({
+  entities: {
+    todos: i.entity({
+      title: i.string(),
+    }),
+    users: i.entity({
+      email: i.string(),
+    }),
+    'user-profiles': i.entity({
+      'display-name': i.string(),
+    }),
+  },
+  links: {
+  },
+  rooms: {},
+});
+
+export default _schema;
+`;
+  const serverSchema = i.schema({
+    entities: {
+      todos: i.entity({ title: i.string() }),
+      users: i.entity({ email: i.string() }),
+      'user-profiles': i.entity({
+        'display-name': i.string(),
+        'avatar-url': i.string(),
+      }),
+    },
+    links: {
+      'todo-owner': {
+        forward: { on: 'todos', has: 'one', label: 'owner' },
+        reverse: { on: 'users', has: 'many', label: 'todos' },
+      },
+    },
+  });
+
+  const result = await update(oldFile, serverSchema);
+
+  expect(result).toMatchSnapshot();
+});
+
+test('adds a link when links object is empty', async () => {
+  const oldFile = `
+import { i } from '@instantdb/core';
+
+const _schema = i.schema({
+  entities: {
+    todos: i.entity({
+      title: i.string(),
+    }),
+    users: i.entity({
+      email: i.string(),
+    }),
+  },
+  links: {
+  },
+  rooms: {},
+});
+
+export default _schema;
+`;
+  const serverSchema = i.schema({
+    entities: {
+      todos: i.entity({ title: i.string() }),
+      users: i.entity({ email: i.string() }),
+    },
+    links: {
+      todoOwner: {
+        forward: { on: 'todos', has: 'one', label: 'owner' },
+        reverse: { on: 'users', has: 'many', label: 'todos' },
+      },
+    },
+  });
+
+  const result = await update(oldFile, serverSchema);
+
+  expect(result).toMatchSnapshot();
+});
+
+test('adds multiple links when links object is empty', async () => {
+  const oldFile = `
+import { i } from '@instantdb/core';
+
+const _schema = i.schema({
+  entities: {
+    todos: i.entity({
+      title: i.string(),
+    }),
+    users: i.entity({
+      email: i.string(),
+    }),
+    projects: i.entity({
+      name: i.string(),
+    }),
+  },
+  links: {
+  },
+  rooms: {},
+});
+
+export default _schema;
+`;
+  const serverSchema = i.schema({
+    entities: {
+      todos: i.entity({ title: i.string() }),
+      users: i.entity({ email: i.string() }),
+      projects: i.entity({ name: i.string() }),
+    },
+    links: {
+      todoOwner: {
+        forward: { on: 'todos', has: 'one', label: 'owner' },
+        reverse: { on: 'users', has: 'many', label: 'todos' },
+      },
+      projectTodos: {
+        forward: { on: 'projects', has: 'many', label: 'todos' },
+        reverse: { on: 'todos', has: 'one', label: 'project' },
+      },
+    },
+  });
+
+  const result = await update(oldFile, serverSchema);
+
+  await expectNoDiff(result, serverSchema);
+});
+
+test('removes the last link cleanly', async () => {
+  const oldFile = `
+import { i } from '@instantdb/core';
+
+const _schema = i.schema({
+  entities: {
+    todos: i.entity({
+      title: i.string(),
+    }),
+    users: i.entity({
+      email: i.string(),
+    }),
+  },
+  links: {
+    todoOwner: {
+      forward: { on: 'todos', has: 'one', label: 'owner' },
+      reverse: { on: 'users', has: 'many', label: 'todos' },
+    },
+  },
+  rooms: {},
+});
+
+export default _schema;
+`;
+  const serverSchema = i.schema({
+    entities: {
+      todos: i.entity({ title: i.string() }),
+      users: i.entity({ email: i.string() }),
+    },
+    links: {},
+  });
+
+  const result = await update(oldFile, serverSchema);
+
+  expect(result).toMatchSnapshot();
+});
+
+async function update(oldFile: string, serverSchema: any) {
+  const localSchema = schemaTypescriptFileToInstantSchema(oldFile);
+  return updateSchemaFile(oldFile, localSchema, serverSchema);
+}
+
+async function expectNoDiff(result: string, serverSchema: any) {
+  const parsed = schemaTypescriptFileToInstantSchema(result);
+  const diff = await diffSchemas(parsed, serverSchema, async (created) => created, {});
+  expect(diff).toEqual([]);
+}
diff --git a/client/packages/cli/src/index.js b/client/packages/cli/src/index.js
index fe92b63413..a5dd73103b 100644
--- a/client/packages/cli/src/index.js
+++ b/client/packages/cli/src/index.js
@@ -10,7 +10,6 @@ import {
   PlatformApi,
 } from '@instantdb/platform';
 import version from './version.js';
-import { existsSync } from 'fs';
 import { mkdir, writeFile, readFile, unlink } from 'fs/promises';
 import path, { join } from 'path';
 import { randomUUID } from 'crypto';
@@ -50,7 +49,8 @@ import {
   getSchemaPathToWrite,
   getPermsPathToWrite,
 } from './util/findConfigCandidates.js';
-import { mergeSchema } from './util/mergeSchema.js';
+
+import { updateSchemaFile } from './util/updateSchemaFile.js';
 
 const execAsync = promisify(exec);
 
@@ -481,7 +481,7 @@ program
   )
   .option(
     '--experimental-type-preservation',
-    "[Experimental] Preserve manual type changes like `status: i.json<'online' | 'offline'>()` when doing `instant-cli pull schema`",
+    '[Experimental] Preserve manual type changes and schema edits when pulling schema',
   )
   .description('Pull schema and perm files from production.')
   .addHelpText(
@@ -1236,24 +1236,33 @@ async function pullSchema(
   const shortSchemaPath = getSchemaPathToWrite(prev?.path);
   const schemaPath = join(pkgDir, shortSchemaPath);
 
-  let newSchemaContent = generateSchemaTypescriptFile(
-    prev?.schema,
-    apiSchemaToInstantSchemaDef(pullRes.data.schema),
-    instantModuleName,
-  );
+  const serverSchema = apiSchemaToInstantSchemaDef(pullRes.data.schema);
+  let newSchemaContent;
 
   if (prev && experimentalTypePreservation) {
     try {
       const oldSchemaContent = await readFile(prev.path, 'utf-8');
-      newSchemaContent = mergeSchema(oldSchemaContent, newSchemaContent);
+      newSchemaContent = await updateSchemaFile(
+        oldSchemaContent,
+        prev.schema,
+        serverSchema,
+      );
     } catch (e) {
       warn(
-        'Failed to merge schema with existing file. Overwriting instead.',
+        'Failed to update schema with existing file. Overwriting instead.',
         e,
       );
     }
   }
 
+  if (!newSchemaContent) {
+    newSchemaContent = generateSchemaTypescriptFile(
+      prev?.schema,
+      serverSchema,
+      instantModuleName,
+    );
+  }
+
   await writeTypescript(schemaPath, newSchemaContent, 'utf-8');
 
   console.log('✅ Wrote schema to ' + shortSchemaPath);
diff --git a/client/packages/cli/src/util/mergeSchema.js b/client/packages/cli/src/util/mergeSchema.js
deleted file mode 100644
index afbf3346c4..0000000000
--- a/client/packages/cli/src/util/mergeSchema.js
+++ /dev/null
@@ -1,364 +0,0 @@
-import * as acorn from 'acorn';
-import tsPlugin from 'acorn-typescript';
-
-const node = acorn.Parser.extend(tsPlugin({ dts: false }));
-
-// --- Import Handling Helpers ---
-
-function collectImports(ast) {
-  const imports = new Map(); // localName -> { source, importedName, type, importKind }
-
-  for (const node of ast.body) {
-    if (node.type === 'ImportDeclaration') {
-      const source = node.source.value;
-      const declImportKind = node.importKind;
-
-      for (const specifier of node.specifiers) {
-        let kind = 'value';
-        if (declImportKind === 'type' || specifier.importKind === 'type') {
-          kind = 'type';
-        }
-
-        if (specifier.type === 'ImportSpecifier') {
-          imports.set(specifier.local.name, {
-            source,
-            importedName: specifier.imported.name,
-            type: 'named',
-            importKind: kind,
-          });
-        } else if (specifier.type === 'ImportDefaultSpecifier') {
-          imports.set(specifier.local.name, {
-            source,
-            type: 'default',
-            importKind: kind,
-          });
-        } else if (specifier.type === 'ImportNamespaceSpecifier') {
-          imports.set(specifier.local.name, {
-            source,
-            type: 'namespace',
-            importKind: kind,
-          });
-        }
-      }
-    }
-  }
-  return imports;
-}
-
-function collectExistingImports(ast) {
-  const existing = new Map(); // source -> Set<localName>
-
-  for (const node of ast.body) {
-    if (node.type === 'ImportDeclaration') {
-      const source = node.source.value;
-      if (!existing.has(source)) existing.set(source, new Set());
-
-      for (const specifier of node.specifiers) {
-        existing.get(source).add(specifier.local.name);
-      }
-    }
-  }
-  return existing;
-}
-
-function findIdentifiers(node, set = new Set()) {
-  if (!node) return set;
-
-  // If it's a Type Reference, grab the name
-  if (node.type === 'TSTypeReference') {
-    if (node.typeName.type === 'Identifier') {
-      set.add(node.typeName.name);
-    } else if (node.typeName.type === 'TSQualifiedName') {
-      // For A.B, we need A
-      let left = node.typeName.left;
-      while (left.type === 'TSQualifiedName') {
-        left = left.left;
-      }
-      if (left.type === 'Identifier') {
-        set.add(left.name);
-      }
-    }
-  }
-
-  // Recursive walk
-  for (const key in node) {
-    if (key === 'loc' || key === 'start' || key === 'end') continue;
-    if (typeof node[key] === 'object' && node[key] !== null) {
-      if (Array.isArray(node[key])) {
-        node[key].forEach((child) => findIdentifiers(child, set));
-      } else {
-        findIdentifiers(node[key], set);
-      }
-    }
-  }
-  return set;
-}
-
-// --- Schema Traversal Helpers ---
-
-function getPropName(prop) {
-  if (prop.key.type === 'Identifier') return prop.key.name;
-  if (prop.key.type === 'Literal') return prop.key.value;
-  return null;
-}
-
-function analyzeChain(node) {
-  let curr = node;
-  let typeParams = null;
-  let baseCall = null;
-
-  while (curr.type === 'CallExpression') {
-    if (curr.typeParameters) {
-      typeParams = curr.typeParameters;
-    }
-
-    if (
-      curr.callee.type === 'MemberExpression' &&
-      curr.callee.object.type === 'Identifier' &&
-      curr.callee.object.name === 'i'
-    ) {
-      baseCall = curr;
-    }
-
-    if (curr.callee.type === 'MemberExpression') {
-      curr = curr.callee.object;
-    } else {
-      break;
-    }
-  }
-  return { typeParams, baseCall };
-}
-
-function traverseSchema(node, path, callback) {
-  if (node.type === 'ObjectExpression') {
-    for (const prop of node.properties) {
-      const name = getPropName(prop);
-      if (!name) continue;
-      const newPath = path ? `${path}.${name}` : name;
-
-      if (prop.value.type === 'ObjectExpression') {
-        traverseSchema(prop.value, newPath, callback);
-      } else if (prop.value.type === 'CallExpression') {
-        let isEntity = false;
-        if (
-          prop.value.callee.type === 'MemberExpression' &&
-          prop.value.callee.property.name === 'entity' &&
-          prop.value.callee.object.type === 'Identifier' &&
-          prop.value.callee.object.name === 'i'
-        ) {
-          isEntity = true;
-        }
-
-        if (isEntity) {
-          callback(prop.value, newPath);
-          if (prop.value.arguments.length > 0) {
-            traverseSchema(prop.value.arguments[0], newPath, callback);
-          }
-        } else {
-          callback(prop.value, newPath);
-        }
-      }
-    }
-  }
-}
-
-function findSchemaObject(ast) {
-  let schemaObj = null;
-  function walk(node) {
-    if (!node) return;
-    if (schemaObj) return;
-    if (
-      node.type === 'CallExpression' &&
-      node.callee.type === 'MemberExpression' &&
-      node.callee.object.name === 'i' &&
-      node.callee.property.name === 'schema' &&
-      node.arguments.length > 0
-    ) {
-      schemaObj = node.arguments[0];
-      return;
-    }
-    for (const key in node) {
-      if (key === 'loc' || key === 'start' || key === 'end') continue;
-      if (node[key] && typeof node[key] === 'object') {
-        if (Array.isArray(node[key])) {
-          node[key].forEach(walk);
-        } else {
-          walk(node[key]);
-        }
-      }
-    }
-  }
-  walk(ast);
-  return schemaObj;
-}
-
-export function mergeSchema(oldFile, newFile) {
-  const oldParsed = node.parse(oldFile, {
-    sourceType: 'module',
-    ecmaVersion: 'latest',
-    locations: true,
-  });
-
-  const newParsed = node.parse(newFile, {
-    sourceType: 'module',
-    ecmaVersion: 'latest',
-    locations: true,
-  });
-
-  const schemaMap = new Map(); // Path -> { src, ast }
-
-  // 1. Extract from old file
-  const oldSchemaObj = findSchemaObject(oldParsed);
-  if (oldSchemaObj) {
-    traverseSchema(oldSchemaObj, '', (node, path) => {
-      const { typeParams } = analyzeChain(node);
-      if (typeParams) {
-        const src = oldFile.slice(typeParams.start, typeParams.end);
-        schemaMap.set(path, { src, ast: typeParams });
-      }
-    });
-  }
-
-  // 2. Collect Imports
-  const oldImports = collectImports(oldParsed);
-  const newExistingImports = collectExistingImports(newParsed);
-  const neededIdentifiers = new Set();
-
-  // 3. Apply to new file & Collect needed identifiers
-  const edits = [];
-  const newSchemaObj = findSchemaObject(newParsed);
-
-  if (newSchemaObj) {
-    traverseSchema(newSchemaObj, '', (node, path) => {
-      const { typeParams, baseCall } = analyzeChain(node);
-      const stored = schemaMap.get(path);
-
-      if (stored) {
-        // Collect identifiers from the type params we are about to inject
-        findIdentifiers(stored.ast, neededIdentifiers);
-
-        if (typeParams) {
-          edits.push({
-            start: typeParams.start,
-            end: typeParams.end,
-            text: stored.src,
-          });
-        } else if (baseCall) {
-          edits.push({
-            start: baseCall.callee.end,
-            end: baseCall.callee.end,
-            text: stored.src,
-          });
-        }
-      }
-    });
-  }
-
-  // 4. Generate Import Statements
-  const importsToAdd = new Map(); // source -> { named: Map<string, {str, isType}>, default: {str, isType}, namespace: {str, isType} }
-
-  for (const id of neededIdentifiers) {
-    const importInfo = oldImports.get(id);
-    if (importInfo) {
-      // Check if already imported in new file
-      const existing = newExistingImports.get(importInfo.source);
-      if (existing && existing.has(id)) {
-        continue; // Already imported
-      }
-
-      if (!importsToAdd.has(importInfo.source)) {
-        importsToAdd.set(importInfo.source, {
-          named: new Map(),
-          default: null,
-          namespace: null,
-        });
-      }
-      const group = importsToAdd.get(importInfo.source);
-      const isType = importInfo.importKind === 'type';
-
-      if (importInfo.type === 'named') {
-        let importStr = id;
-        if (importInfo.importedName !== id) {
-          importStr = `${importInfo.importedName} as ${id}`;
-        }
-        group.named.set(importStr, { str: importStr, isType });
-      } else if (importInfo.type === 'default') {
-        group.default = { str: id, isType };
-      } else if (importInfo.type === 'namespace') {
-        group.namespace = { str: id, isType };
-      }
-    }
-  }
-
-  const importBlocks = [];
-
-  for (const [source, info] of importsToAdd) {
-    // Check if source exists in new file to merge?
-    // For simplicity, we append new import lines.
-    // But we can try to be smart.
-
-    // If we have named imports
-    if (info.named.size > 0) {
-      const namedImports = Array.from(info.named.values());
-      const allTypes = namedImports.every((x) => x.isType);
-
-      if (allTypes) {
-        const names = namedImports.map((x) => x.str).join(', ');
-        importBlocks.push(`import type { ${names} } from '${source}';`);
-      } else {
-        const names = namedImports
-          .map((x) => (x.isType ? `type ${x.str}` : x.str))
-          .join(', ');
-        importBlocks.push(`import { ${names} } from '${source}';`);
-      }
-    }
-    if (info.default) {
-      if (info.default.isType) {
-        importBlocks.push(`import type ${info.default.str} from '${source}';`);
-      } else {
-        importBlocks.push(`import ${info.default.str} from '${source}';`);
-      }
-    }
-    if (info.namespace) {
-      if (info.namespace.isType) {
-        importBlocks.push(
-          `import type * as ${info.namespace.str} from '${source}';`,
-        );
-      } else {
-        importBlocks.push(
-          `import * as ${info.namespace.str} from '${source}';`,
-        );
-      }
-    }
-  }
-
-  // 5. Apply edits
-  edits.sort((a, b) => b.start - a.start);
-
-  let output = newFile;
-  for (const edit of edits) {
-    output = output.slice(0, edit.start) + edit.text + output.slice(edit.end);
-  }
-
-  // Prepend imports
-  if (importBlocks.length > 0) {
-    // Check for leading comments (e.g. // Docs: ...)
-    // We want to insert imports AFTER the leading comments but BEFORE the first code/import
-
-    // Simple heuristic: match consecutive lines at start that begin with //
-    const commentMatch = output.match(/^(\s*\/\/.*(\r?\n|$))+/);
-
-    if (commentMatch) {
-      const commentEnd = commentMatch[0].length;
-      output =
-        output.slice(0, commentEnd) +
-        '\n' +
-        importBlocks.join('\n') +
-        output.slice(commentEnd);
-    } else {
-      output = importBlocks.join('\n') + '\n' + output;
-    }
-  }
-
-  return output;
-}
diff --git a/client/packages/cli/src/util/updateSchemaFile.ts b/client/packages/cli/src/util/updateSchemaFile.ts
new file mode 100644
index 0000000000..2d8849d308
--- /dev/null
+++ b/client/packages/cli/src/util/updateSchemaFile.ts
@@ -0,0 +1,760 @@
+import * as acorn from 'acorn';
+import { tsPlugin } from 'acorn-typescript';
+import {
+  diffSchemas,
+  renderAttrCall,
+  renderAttrProperty,
+  renderEntityProperty,
+  renderLinkProperty,
+  renderLinkValue,
+  type MigrationTx,
+} from '@instantdb/platform';
+import type { DataAttrDef, InstantSchemaDef } from '@instantdb/core';
+
+const parser = (acorn.Parser as any).extend(tsPlugin({ dts: false }) as any);
+const DEFAULT_INDENT = '  ';
+
+export async function updateSchemaFile(
+  existingFileContent: string,
+  localSchema: InstantSchemaDef<any, any, any>,
+  serverSchema: InstantSchemaDef<any, any, any>,
+): Promise<string> {
+  const ast = parseFile(existingFileContent);
+  const schemaObj = findSchemaObject(ast);
+  if (!schemaObj) {
+    throw new Error('Could not find i.schema(...) in schema file.');
+  }
+
+  const { entitiesObj, linksObj } = getSchemaSections(schemaObj);
+  const diff = await diffSchemas(
+    localSchema,
+    serverSchema,
+    async (created) => created,
+    {},
+  );
+  if (diff.length === 0) {
+    return existingFileContent;
+  }
+
+  const { entitiesByName } = buildEntitiesIndex(entitiesObj);
+
+  const { localLinksByForward, serverLinksByForward } = buildLinkMaps(
+    localSchema.links || {},
+    serverSchema.links || {},
+  );
+  const changeBuckets = collectChanges(
+    diff,
+    localLinksByForward,
+    serverLinksByForward,
+  );
+
+  const edits: Edit[] = [];
+  edits.push(
+    ...collectEntityEdits(
+      existingFileContent,
+      entitiesObj,
+      entitiesByName,
+      changeBuckets,
+      serverSchema,
+    ),
+  );
+  edits.push(
+    ...collectLinkEdits(
+      existingFileContent,
+      linksObj,
+      changeBuckets,
+      localLinksByForward,
+      serverLinksByForward,
+    ),
+  );
+
+  const updated = applyEdits(existingFileContent, edits);
+  return normalizeEmptyLinksObject(updated);
+}
+
+type ObjectExpression = {
+  type: 'ObjectExpression';
+  start: number;
+  end: number;
+  properties: any[];
+};
+
+type PropertyNode = {
+  type: 'Property';
+  start: number;
+  end: number;
+  key: any;
+  value: any;
+};
+
+type EntityInfo = {
+  prop: PropertyNode;
+  attrsObj: ObjectExpression;
+  attrsByName: Map<string, PropertyNode>;
+};
+
+type Edit = { start: number; end: number; text: string };
+
+type ChangeBuckets = {
+  createEntities: Set<string>;
+  deleteEntities: Set<string>;
+  addAttrs: Map<string, Set<string>>;
+  deleteAttrs: Map<string, Set<string>>;
+  updateAttrs: Map<string, Set<string>>;
+  addLinks: Set<string>;
+  deleteLinks: Set<string>;
+  updateLinks: Set<string>;
+};
+
+type LinkMap = Map<string, { name: string; link: any }>;
+
+function parseFile(content: string) {
+  return parser.parse(content, {
+    sourceType: 'module',
+    ecmaVersion: 'latest',
+  });
+}
+
+function getSchemaSections(schemaObj: ObjectExpression) {
+  const entitiesProp = findObjectProperty(schemaObj, 'entities');
+  if (!entitiesProp || !isObjectExpression(entitiesProp.value)) {
+    throw new Error('Could not find entities object in schema file.');
+  }
+  const linksProp = findObjectProperty(schemaObj, 'links');
+  if (!linksProp || !isObjectExpression(linksProp.value)) {
+    throw new Error('Could not find links object in schema file.');
+  }
+
+  return {
+    entitiesObj: entitiesProp.value,
+    linksObj: linksProp.value,
+  };
+}
+
+function buildEntitiesIndex(entitiesObj: ObjectExpression) {
+  const entitiesByName = new Map<string, EntityInfo>();
+
+  for (const prop of entitiesObj.properties) {
+    if (!isProperty(prop)) continue;
+    const name = getPropName(prop);
+    if (!name) continue;
+    const attrsObj = getEntityAttrsObject(prop.value);
+    if (!attrsObj) continue;
+
+    const attrsByName = new Map<string, PropertyNode>();
+    for (const attrProp of attrsObj.properties) {
+      if (!isProperty(attrProp)) continue;
+      const attrName = getPropName(attrProp);
+      if (!attrName) continue;
+      attrsByName.set(attrName, attrProp);
+    }
+
+    entitiesByName.set(name, { prop, attrsObj, attrsByName });
+  }
+
+  return { entitiesByName };
+}
+
+function buildLinkMaps(
+  localLinks: Record<string, any>,
+  serverLinks: Record<string, any>,
+) {
+  return {
+    localLinksByForward: buildLinkForwardMap(localLinks),
+    serverLinksByForward: buildLinkForwardMap(serverLinks),
+  };
+}
+
+function collectEntityEdits(
+  content: string,
+  entitiesObj: ObjectExpression,
+  entitiesByName: Map<string, EntityInfo>,
+  changeBuckets: ChangeBuckets,
+  serverSchema: InstantSchemaDef<any, any, any>,
+): Edit[] {
+  const edits: Edit[] = [];
+
+  for (const entityName of changeBuckets.deleteEntities) {
+    const entity = entitiesByName.get(entityName);
+    if (!entity) continue;
+    edits.push(removeProperty(content, entitiesObj, entity.prop));
+  }
+
+  for (const entityName of changeBuckets.createEntities) {
+    if (entitiesByName.has(entityName)) continue;
+    const entityDef = serverSchema.entities?.[entityName];
+    if (!entityDef) continue;
+    const propText = renderEntityProperty(entityName, entityDef.attrs);
+    const indent = getObjectPropIndent(content, entitiesObj);
+    edits.push(insertProperty(content, entitiesObj, propText, indent));
+  }
+
+  for (const [entityName, attrs] of changeBuckets.deleteAttrs) {
+    const entity = entitiesByName.get(entityName);
+    if (!entity) continue;
+    for (const attrName of attrs) {
+      const attrProp = entity.attrsByName.get(attrName);
+      if (!attrProp) continue;
+      edits.push(removeProperty(content, entity.attrsObj, attrProp));
+    }
+  }
+
+  for (const [entityName, attrs] of changeBuckets.addAttrs) {
+    const entity = entitiesByName.get(entityName);
+    if (!entity) continue;
+    const entityDef = serverSchema.entities?.[entityName];
+    if (!entityDef) continue;
+    const indent = getObjectPropIndent(content, entity.attrsObj);
+    for (const attrName of attrs) {
+      const attrDef = entityDef.attrs?.[attrName];
+      if (!attrDef) continue;
+      const propText = renderAttrProperty(attrName, attrDef);
+      edits.push(
+        insertPropertyExpandingSingleLine(
+          content,
+          entity.attrsObj,
+          propText,
+          indent,
+        ),
+      );
+    }
+  }
+
+  for (const [entityName, attrs] of changeBuckets.updateAttrs) {
+    const entity = entitiesByName.get(entityName);
+    if (!entity) continue;
+    const entityDef = serverSchema.entities?.[entityName];
+    if (!entityDef) continue;
+    for (const attrName of attrs) {
+      const attrProp = entity.attrsByName.get(attrName);
+      const attrDef = entityDef.attrs?.[attrName];
+      if (!attrProp || !attrDef) continue;
+      edits.push(updateAttrEdit(content, attrProp, attrDef));
+    }
+  }
+
+  return edits;
+}
+
+function collectLinkEdits(
+  content: string,
+  linksObj: ObjectExpression,
+  changeBuckets: ChangeBuckets,
+  localLinksByForward: LinkMap,
+  serverLinksByForward: LinkMap,
+): Edit[] {
+  const edits: Edit[] = [];
+
+  for (const forwardKey of changeBuckets.deleteLinks) {
+    const localLink = localLinksByForward.get(forwardKey);
+    if (!localLink) continue;
+    const linkProp = findObjectProperty(linksObj, localLink.name);
+    if (!linkProp) continue;
+    edits.push(removeProperty(content, linksObj, linkProp));
+  }
+
+  for (const forwardKey of changeBuckets.addLinks) {
+    const serverLink = serverLinksByForward.get(forwardKey);
+    if (!serverLink) continue;
+    const linkName = serverLink.name;
+    if (findObjectProperty(linksObj, linkName)) continue;
+    const propText = renderLinkProperty(linkName, serverLink.link);
+    const indent = getObjectPropIndent(content, linksObj);
+    edits.push(insertProperty(content, linksObj, propText, indent));
+  }
+
+  for (const forwardKey of changeBuckets.updateLinks) {
+    const serverLink = serverLinksByForward.get(forwardKey);
+    const localLink = localLinksByForward.get(forwardKey);
+    if (!serverLink || !localLink) continue;
+    const linkProp = findObjectProperty(linksObj, localLink.name);
+    if (!linkProp) continue;
+    const nextValue = renderLinkValue(serverLink.link);
+    const propIndent = getLineIndent(content, linkProp.start);
+    edits.push({
+      start: linkProp.value.start,
+      end: linkProp.value.end,
+      text: indentValueAfterFirstLine(nextValue, propIndent),
+    });
+  }
+
+  return edits;
+}
+
+function updateAttrEdit(
+  content: string,
+  attrProp: PropertyNode,
+  attrDef: DataAttrDef<string, boolean, boolean>,
+): Edit {
+  const { typeParams } = analyzeChain(attrProp.value);
+  const typeParamsText = typeParams
+    ? content.slice(typeParams.start, typeParams.end)
+    : null;
+  const nextValue = renderAttrCall(attrDef, typeParamsText);
+  return {
+    start: attrProp.value.start,
+    end: attrProp.value.end,
+    text: nextValue,
+  };
+}
+
+function applyEdits(content: string, edits: Edit[]) {
+  if (edits.length === 0) {
+    return content;
+  }
+  const sorted = [...edits].sort((a, b) => b.start - a.start);
+  let output = content;
+  for (const edit of sorted) {
+    output = output.slice(0, edit.start) + edit.text + output.slice(edit.end);
+  }
+  return output;
+}
+
+function collectChanges(
+  steps: MigrationTx[],
+  localLinksByForward: LinkMap,
+  serverLinksByForward: LinkMap,
+): ChangeBuckets {
+  const buckets: ChangeBuckets = {
+    createEntities: new Set(),
+    deleteEntities: new Set(),
+    addAttrs: new Map(),
+    deleteAttrs: new Map(),
+    updateAttrs: new Map(),
+    addLinks: new Set(),
+    deleteLinks: new Set(),
+    updateLinks: new Set(),
+  };
+
+  for (const step of steps) {
+    const namespace = step.identifier.namespace;
+    const attrName = step.identifier.attrName;
+    const forwardKey = `${namespace}.${attrName}`;
+
+    switch (step.type) {
+      case 'add-attr': {
+        if (attrName === 'id') {
+          buckets.createEntities.add(namespace);
+          break;
+        }
+        if (step['value-type'] === 'ref') {
+          buckets.addLinks.add(forwardKey);
+          break;
+        }
+        ensureSet(buckets.addAttrs, namespace).add(attrName);
+        break;
+      }
+      case 'delete-attr': {
+        if (attrName === 'id') {
+          buckets.deleteEntities.add(namespace);
+          break;
+        }
+        if (localLinksByForward.has(forwardKey)) {
+          buckets.deleteLinks.add(forwardKey);
+          break;
+        }
+        ensureSet(buckets.deleteAttrs, namespace).add(attrName);
+        break;
+      }
+      case 'update-attr': {
+        if (step.partialAttr?.['value-type'] === 'ref') {
+          buckets.updateLinks.add(forwardKey);
+          break;
+        }
+        ensureSet(buckets.updateAttrs, namespace).add(attrName);
+        break;
+      }
+      case 'index':
+      case 'remove-index':
+      case 'unique':
+      case 'remove-unique':
+      case 'required':
+      case 'remove-required':
+      case 'check-data-type':
+      case 'remove-data-type': {
+        if (serverLinksByForward.has(forwardKey)) {
+          buckets.updateLinks.add(forwardKey);
+          break;
+        }
+        ensureSet(buckets.updateAttrs, namespace).add(attrName);
+        break;
+      }
+      default: {
+        assertNever(step);
+      }
+    }
+  }
+
+  pruneNamespaceBuckets(buckets);
+  return buckets;
+}
+
+function ensureSet(map: Map<string, Set<string>>, key: string) {
+  if (!map.has(key)) map.set(key, new Set());
+  return map.get(key)!;
+}
+
+function pruneNamespaceBuckets(buckets: ChangeBuckets) {
+  const namespaces = new Set<string>([
+    ...buckets.createEntities,
+    ...buckets.deleteEntities,
+  ]);
+
+  for (const namespace of namespaces) {
+    buckets.addAttrs.delete(namespace);
+    buckets.deleteAttrs.delete(namespace);
+    buckets.updateAttrs.delete(namespace);
+    removeNamespaceLinks(buckets.addLinks, namespace);
+    removeNamespaceLinks(buckets.deleteLinks, namespace);
+    removeNamespaceLinks(buckets.updateLinks, namespace);
+  }
+}
+
+function removeNamespaceLinks(set: Set<string>, namespace: string) {
+  const prefix = `${namespace}.`;
+  for (const key of Array.from(set)) {
+    if (key.startsWith(prefix)) {
+      set.delete(key);
+    }
+  }
+}
+
+function assertNever(value: never): never {
+  throw new Error(`Unhandled migration step: ${JSON.stringify(value)}`);
+}
+
+function analyzeChain(node: any): { typeParams: acorn.Node | null } {
+  let curr = node;
+  let typeParams: acorn.Node | null = null;
+
+  while (curr?.type === 'CallExpression') {
+    if (curr.typeParameters) {
+      typeParams = curr.typeParameters;
+    }
+    if (curr.callee?.type === 'MemberExpression') {
+      curr = curr.callee.object;
+    } else {
+      break;
+    }
+  }
+  return { typeParams };
+}
+
+function findSchemaObject(ast: any): ObjectExpression | null {
+  let schemaObj: ObjectExpression | null = null;
+  const walk = (node: any) => {
+    if (!node || schemaObj) return;
+    if (
+      node.type === 'CallExpression' &&
+      node.callee?.type === 'MemberExpression' &&
+      node.callee.object?.type === 'Identifier' &&
+      node.callee.object.name === 'i' &&
+      node.callee.property?.name === 'schema' &&
+      node.arguments?.length > 0
+    ) {
+      const arg = node.arguments[0];
+      if (isObjectExpression(arg)) {
+        schemaObj = arg;
+      }
+      return;
+    }
+    for (const key in node) {
+      if (key === 'loc' || key === 'start' || key === 'end') continue;
+      const value = node[key];
+      if (!value || typeof value !== 'object') continue;
+      if (Array.isArray(value)) {
+        value.forEach(walk);
+      } else {
+        walk(value);
+      }
+    }
+  };
+  walk(ast);
+  return schemaObj;
+}
+
+function findObjectProperty(
+  obj: ObjectExpression,
+  name: string,
+): PropertyNode | null {
+  for (const prop of obj.properties) {
+    if (!isProperty(prop)) continue;
+    const propName = getPropName(prop);
+    if (propName === name) return prop;
+  }
+  return null;
+}
+
+function getEntityAttrsObject(value: any): ObjectExpression | null {
+  if (
+    value?.type !== 'CallExpression' ||
+    value.callee?.type !== 'MemberExpression' ||
+    value.callee.object?.type !== 'Identifier' ||
+    value.callee.object.name !== 'i' ||
+    value.callee.property?.name !== 'entity'
+  ) {
+    return null;
+  }
+  const attrsObj = value.arguments?.[0];
+  return isObjectExpression(attrsObj) ? attrsObj : null;
+}
+
+function buildLinkForwardMap(links: Record<string, any>) {
+  const map = new Map<string, { name: string; link: any }>();
+  for (const [name, link] of Object.entries(links || {})) {
+    map.set(linkForwardKey(link), { name, link });
+  }
+  return map;
+}
+
+function linkForwardKey(link: any) {
+  return `${link.forward.on}.${link.forward.label}`;
+}
+
+function isObjectExpression(node: any): node is ObjectExpression {
+  return node?.type === 'ObjectExpression';
+}
+
+function isProperty(node: any): node is PropertyNode {
+  return node?.type === 'Property';
+}
+
+function getPropName(prop: PropertyNode) {
+  if (prop.key.type === 'Identifier') return prop.key.name;
+  if (prop.key.type === 'Literal') return String(prop.key.value);
+  return null;
+}
+
+function indentLines(text: string, indent: string) {
+  return text
+    .split('\n')
+    .map((line) => (line.length ? indent + line : line))
+    .join('\n');
+}
+
+function getObjectPropIndent(source: string, obj: ObjectExpression) {
+  const props = obj.properties.filter(isProperty);
+  if (props.length > 0) {
+    return getLineIndent(source, props[0].start);
+  }
+  const closingIndent = getLineIndent(source, obj.end - 1);
+  return closingIndent + DEFAULT_INDENT;
+}
+
+function getLineIndent(source: string, pos: number) {
+  const lineStart = source.lastIndexOf('\n', pos - 1) + 1;
+  const match = source.slice(lineStart, pos).match(/^[\t ]*/);
+  return match ? match[0] : '';
+}
+
+function insertProperty(
+  source: string,
+  obj: ObjectExpression,
+  propText: string,
+  indent: string,
+) {
+  const props = obj.properties.filter(isProperty);
+  const closingBrace = obj.end - 1;
+  const propTextWithIndent = indentLines(propText, indent);
+  const propTextSingleLine = propText.trim();
+  const innerStart = obj.start + 1;
+  const innerEnd = closingBrace;
+  const innerContent = source.slice(innerStart, innerEnd);
+  const innerWhitespaceOnly = /^[\s]*$/.test(innerContent);
+
+  if (props.length === 0) {
+    const objSource = source.slice(obj.start, obj.end);
+    const multiline = objSource.includes('\n') || propText.includes('\n');
+    if (!multiline) {
+      return {
+        start: closingBrace,
+        end: closingBrace,
+        text: ` ${propTextSingleLine} `,
+      };
+    }
+    const closingIndent = getLineIndent(source, closingBrace);
+    if (innerWhitespaceOnly) {
+      return {
+        start: innerStart,
+        end: innerEnd,
+        text: `\n${propTextWithIndent},\n${closingIndent}`,
+      };
+    }
+    return {
+      start: closingBrace,
+      end: closingBrace,
+      text: `\n${propTextWithIndent},\n${closingIndent}`,
+    };
+  }
+
+  const lastProp = props[props.length - 1];
+  const multiline =
+    source.slice(lastProp.end, closingBrace).includes('\n') ||
+    propText.includes('\n');
+  const needsComma = !hasTrailingComma(source, lastProp.end, obj.end);
+
+  if (!multiline) {
+    let insertPos = closingBrace;
+    while (insertPos > lastProp.end && /\s/.test(source[insertPos - 1])) {
+      insertPos -= 1;
+    }
+    return {
+      start: insertPos,
+      end: insertPos,
+      text: `${needsComma ? ',' : ''} ${propTextSingleLine}`,
+    };
+  }
+
+  const lineStart = source.lastIndexOf('\n', closingBrace);
+  return {
+    start: lineStart,
+    end: lineStart,
+    text: `${needsComma ? ',' : ''}\n${propTextWithIndent},`,
+  };
+}
+
+function removeProperty(
+  source: string,
+  obj: ObjectExpression,
+  prop: PropertyNode,
+) {
+  let start = prop.start;
+  let end = prop.end;
+  const lineStart = source.lastIndexOf('\n', start - 1) + 1;
+  let shouldTrimLineEnd = false;
+  if (/^[\t ]*$/.test(source.slice(lineStart, start))) {
+    start = lineStart;
+    shouldTrimLineEnd = true;
+  }
+  const after = skipWhitespaceAndComments(source, end, obj.end);
+  if (source[after] === ',') {
+    end = after + 1;
+  } else {
+    const before = skipWhitespaceAndCommentsBackward(source, start, obj.start);
+    if (source[before] === ',') {
+      start = before;
+    }
+  }
+  if (shouldTrimLineEnd) {
+    const lineEnd = source.indexOf('\n', end);
+    if (lineEnd !== -1 && /^[\t ]*$/.test(source.slice(end, lineEnd))) {
+      end = lineEnd + 1;
+    }
+  }
+  return { start, end, text: '' };
+}
+
+function indentValueAfterFirstLine(value: string, indent: string) {
+  const lines = value.split('\n');
+  if (lines.length <= 1) return value;
+  return [
+    lines[0],
+    ...lines.slice(1).map((line) => (line ? indent + line : line)),
+  ].join('\n');
+}
+
+function insertPropertyExpandingSingleLine(
+  source: string,
+  obj: ObjectExpression,
+  propText: string,
+  indent: string,
+) {
+  const props = obj.properties.filter(isProperty);
+  if (!props.length) {
+    return insertProperty(source, obj, propText, indent);
+  }
+  const objSource = source.slice(obj.start, obj.end);
+  if (objSource.includes('\n') || propText.includes('\n')) {
+    return insertProperty(source, obj, propText, indent);
+  }
+  const closingBrace = obj.end - 1;
+  const closingIndent = getLineIndent(source, closingBrace);
+  const innerIndent = closingIndent + DEFAULT_INDENT;
+  const propTexts = props.map((prop) =>
+    source.slice(prop.start, prop.end).trim(),
+  );
+  const lines = [...propTexts, propText.trim()].map(
+    (prop) => `${innerIndent}${prop},`,
+  );
+  const nextObject = `{\n${lines.join('\n')}\n${closingIndent}}`;
+  return { start: obj.start, end: obj.end, text: nextObject };
+}
+
+function normalizeEmptyLinksObject(content: string) {
+  let ast: any;
+  try {
+    ast = parseFile(content);
+  } catch {
+    return content;
+  }
+  const schemaObj = findSchemaObject(ast);
+  if (!schemaObj) return content;
+  const linksProp = findObjectProperty(schemaObj, 'links');
+  if (!linksProp || !isObjectExpression(linksProp.value)) return content;
+  if (linksProp.value.properties.some(isProperty)) return content;
+  const inner = content.slice(
+    linksProp.value.start + 1,
+    linksProp.value.end - 1,
+  );
+  if (!/^[\s]*$/.test(inner)) return content;
+  return (
+    content.slice(0, linksProp.value.start) +
+    '{}' +
+    content.slice(linksProp.value.end)
+  );
+}
+
+function hasTrailingComma(source: string, afterPos: number, endPos: number) {
+  const next = skipWhitespaceAndComments(source, afterPos, endPos);
+  return source[next] === ',';
+}
+
+function skipWhitespaceAndComments(source: string, start: number, end: number) {
+  let i = start;
+  while (i < end) {
+    const ch = source[i];
+    if (/\s/.test(ch)) {
+      i += 1;
+      continue;
+    }
+    if (ch === '/' && source[i + 1] === '/') {
+      const nextLine = source.indexOf('\n', i + 2);
+      i = nextLine === -1 ? end : nextLine + 1;
+      continue;
+    }
+    if (ch === '/' && source[i + 1] === '*') {
+      const close = source.indexOf('*/', i + 2);
+      i = close === -1 ? end : close + 2;
+      continue;
+    }
+    break;
+  }
+  return i;
+}
+
+function skipWhitespaceAndCommentsBackward(
+  source: string,
+  start: number,
+  end: number,
+) {
+  let i = start - 1;
+  while (i >= end) {
+    const ch = source[i];
+    if (/\s/.test(ch)) {
+      i -= 1;
+      continue;
+    }
+    if (ch === '/' && source[i - 1] === '/') {
+      const prevLine = source.lastIndexOf('\n', i - 2);
+      i = prevLine === -1 ? end - 1 : prevLine - 1;
+      continue;
+    }
+    if (ch === '/' && source[i - 1] === '*') {
+      const open = source.lastIndexOf('/*', i - 2);
+      i = open === -1 ? end - 1 : open - 1;
+      continue;
+    }
+    break;
+  }
+  return i;
+}
diff --git a/client/packages/platform/src/index.ts b/client/packages/platform/src/index.ts
index 194d8ae6c9..4c5d4c60c6 100644
--- a/client/packages/platform/src/index.ts
+++ b/client/packages/platform/src/index.ts
@@ -56,3 +56,13 @@ export {
   type MigrationTxTypes,
   type Identifier,
 } from './migrations.ts';
+
+export {
+  formatKey,
+  renderAttrCall,
+  renderAttrProperty,
+  renderEntityProperty,
+  renderLinkValue,
+  renderLinkProperty,
+  renderLinksObject,
+} from './schemaCodegen.ts';
diff --git a/client/packages/platform/src/schema.ts b/client/packages/platform/src/schema.ts
index 4d0868d934..b2027e1947 100644
--- a/client/packages/platform/src/schema.ts
+++ b/client/packages/platform/src/schema.ts
@@ -15,9 +15,13 @@ import {
   indentLines,
   joinWithTrailingSep,
   sortedEntries,
-  formatKey,
   GenericSchemaDef,
 } from './util.ts';
+import {
+  formatKey,
+  renderEntityProperty,
+  renderLinksObject,
+} from './schemaCodegen.ts';
 
 export type InstantAPIPlatformSchema = {
   refs: Record<string, InstantDBAttr>;
@@ -258,18 +262,6 @@ export type InstantAPISchemaPushStep =
   | InstantAPISchemaPushCheckDataTypeStep
   | InstantAPISchemaPushRemoveDataTypeStep;
 
-function attrDefToCodeString([name, attr]: [
-  string,
-  DataAttrDef<string, boolean, boolean>,
-]) {
-  const type =
-    (attr.metadata.derivedType as any)?.type || attr.valueType || 'any';
-  const unique = attr.config.unique ? '.unique()' : '';
-  const index = attr.config.indexed ? '.indexed()' : '';
-  const required = attr.required ? '' : '.optional()';
-  return `${formatKey(name)}: i.${type}()${unique}${index}${required}`;
-}
-
 function entityDefToCodeStr(
   name: string,
   edef: EntityDef<
@@ -278,14 +270,10 @@ function entityDefToCodeStr(
     any
   >,
 ) {
-  const attrBlock = joinWithTrailingSep(
-    sortedEntries(edef.attrs).map(attrDefToCodeString),
-    ',\n',
-    ',',
-  );
-
-  // a block of code for each entity
-  return `${formatKey(name)}: i.entity({${attrBlock.length ? '\n' : ''}${indentLines(attrBlock, 2)}${attrBlock.length ? '\n' : ''}})`;
+  return renderEntityProperty(name, edef.attrs, {
+    keyOptions: { alwaysQuote: true, quote: '"' },
+    trailingComma: true,
+  });
 }
 
 export function identEtype(ident: InstantDBIdent) {
@@ -338,7 +326,7 @@ function roomDefToCodeStr(room: RoomsDef[string]) {
   }
 
   if (room.topics) {
-    ret += `\n    "topics": {`;
+    ret += `\n    ${formatKey('topics', { alwaysQuote: true, quote: '"' })}: {`;
 
     for (const [topicName, topicConfig] of Object.entries(room.topics)) {
       ret += `\n${indentLines(entityDefToCodeStr(topicName, topicConfig), 6)},`;
@@ -355,7 +343,7 @@ function roomsCodeStr(rooms: RoomsDef) {
   let ret = '{';
 
   for (const [roomType, roomDef] of Object.entries(rooms)) {
-    ret += `\n  ${formatKey(roomType)}: ${roomDefToCodeStr(roomDef)},`;
+    ret += `\n  ${formatKey(roomType, { alwaysQuote: true, quote: '"' })}: ${roomDefToCodeStr(roomDef)},`;
   }
   ret += ret === '{' ? '}' : '\n}';
 
@@ -395,7 +383,11 @@ export function generateSchemaTypescriptFile(
 // run \`push schema\` again to enforce the types.`
       : '';
 
-  const linksEntriesCode = JSON.stringify(newSchema.links, null, 2).trim();
+  const linksEntriesCode = renderLinksObject(newSchema.links, {
+    keyOptions: { alwaysQuote: true, quote: '"' },
+    stringQuote: '"',
+    style: 'json',
+  });
 
   // rooms
   const rooms = prevSchema?.rooms || {};
diff --git a/client/packages/platform/src/schemaCodegen.ts b/client/packages/platform/src/schemaCodegen.ts
new file mode 100644
index 0000000000..6687ba3ea4
--- /dev/null
+++ b/client/packages/platform/src/schemaCodegen.ts
@@ -0,0 +1,183 @@
+import type {
+  AttrsDefs,
+  DataAttrDef,
+  LinkDef,
+  LinksDef,
+} from '@instantdb/core';
+import { indentLines, joinWithTrailingSep, sortedEntries } from './util.ts';
+
+const DEFAULT_INDENT = 2;
+
+type AnyLink = LinkDef<any, any, any, any, any, any, any>;
+
+type KeyFormatOptions = {
+  alwaysQuote?: boolean;
+  quote?: '"' | "'";
+};
+
+type EntityRenderOptions = {
+  indent?: number;
+  typeParamsByAttr?: Record<string, string | null | undefined>;
+  keyOptions?: KeyFormatOptions;
+  trailingComma?: boolean;
+};
+
+type LinkRenderOptions = {
+  indent?: number;
+  keyOptions?: KeyFormatOptions;
+  stringQuote?: '"' | "'";
+  style?: 'js' | 'json';
+};
+
+type LinksObjectOptions = {
+  indent?: number;
+  newlineForEmpty?: boolean;
+  keyOptions?: KeyFormatOptions;
+  stringQuote?: '"' | "'";
+  style?: 'js' | 'json';
+};
+
+export function formatKey(name: string, options: KeyFormatOptions = {}) {
+  const quote = options.quote ?? "'";
+  if (!options.alwaysQuote && isValidIdentifier(name)) {
+    return name;
+  }
+  return `${quote}${name}${quote}`;
+}
+
+export function renderAttrCall(
+  attr: DataAttrDef<string, boolean, boolean>,
+  typeParams?: string | null,
+) {
+  const type =
+    (attr.metadata?.derivedType as any)?.type || attr.valueType || 'any';
+  const unique = attr.config.unique ? '.unique()' : '';
+  const index = attr.config.indexed ? '.indexed()' : '';
+  const required = attr.required ? '' : '.optional()';
+  return `i.${type}${typeParams ?? ''}()${unique}${index}${required}`;
+}
+
+export function renderAttrProperty(
+  name: string,
+  attr: DataAttrDef<string, boolean, boolean>,
+  options: { typeParams?: string | null; keyOptions?: KeyFormatOptions } = {},
+) {
+  return `${formatKey(name, options.keyOptions)}: ${renderAttrCall(
+    attr,
+    options.typeParams,
+  )}`;
+}
+
+export function renderEntityProperty(
+  name: string,
+  attrs: AttrsDefs,
+  options: EntityRenderOptions = {},
+) {
+  const indent = options.indent ?? DEFAULT_INDENT;
+  const keyOptions = options.keyOptions;
+  const attrBlock = sortedEntries(attrs)
+    .map(([attrName, attrDef]) =>
+      renderAttrProperty(attrName, attrDef, {
+        typeParams: options.typeParamsByAttr?.[attrName] ?? null,
+        keyOptions,
+      }),
+    )
+    .filter(Boolean);
+  const attrBlockText = options.trailingComma
+    ? joinWithTrailingSep(attrBlock, ',\n', ',')
+    : attrBlock.join(',\n');
+  const inner = attrBlockText.length
+    ? `\n${indentLines(attrBlockText, indent)}\n`
+    : '';
+  return `${formatKey(name, keyOptions)}: i.entity({${inner}})`;
+}
+
+export function renderLinkValue(
+  link: AnyLink,
+  options: LinkRenderOptions = {},
+) {
+  const indent = options.indent ?? DEFAULT_INDENT;
+  const indentStr = ' '.repeat(indent);
+  const keyOptions = options.keyOptions;
+  const q = options.stringQuote ?? "'";
+  const style = options.style ?? 'js';
+  const forwardLines = [
+    `${formatKey('on', keyOptions)}: ${q}${link.forward.on}${q}`,
+    `${formatKey('has', keyOptions)}: ${q}${link.forward.has}${q}`,
+    `${formatKey('label', keyOptions)}: ${q}${link.forward.label}${q}`,
+  ];
+  if (link.forward.required) {
+    forwardLines.push(`${formatKey('required', keyOptions)}: true`);
+  }
+  if (link.forward.onDelete) {
+    forwardLines.push(
+      `${formatKey('onDelete', keyOptions)}: ${q}${link.forward.onDelete}${q}`,
+    );
+  }
+
+  const reverseLines = [
+    `${formatKey('on', keyOptions)}: ${q}${link.reverse.on}${q}`,
+    `${formatKey('has', keyOptions)}: ${q}${link.reverse.has}${q}`,
+    `${formatKey('label', keyOptions)}: ${q}${link.reverse.label}${q}`,
+  ];
+  if (link.reverse.onDelete) {
+    reverseLines.push(
+      `${formatKey('onDelete', keyOptions)}: ${q}${link.reverse.onDelete}${q}`,
+    );
+  }
+
+  const forwardBodyLines =
+    style === 'json' ? forwardLines : forwardLines.map((line) => `${line},`);
+  const reverseBodyLines =
+    style === 'json' ? reverseLines : reverseLines.map((line) => `${line},`);
+
+  return [
+    '{',
+    `${indentStr}${formatKey('forward', keyOptions)}: {`,
+    `${indentStr}${indentStr}${forwardBodyLines.join(
+      `${style === 'json' ? ',' : ''}\n${indentStr}${indentStr}`,
+    )}`,
+    `${indentStr}}${style === 'json' ? ',' : ','}`,
+    `${indentStr}${formatKey('reverse', keyOptions)}: {`,
+    `${indentStr}${indentStr}${reverseBodyLines.join(
+      `${style === 'json' ? ',' : ''}\n${indentStr}${indentStr}`,
+    )}`,
+    `${indentStr}}${style === 'json' ? '' : ','}`,
+    '}',
+  ].join('\n');
+}
+
+export function renderLinkProperty(
+  name: string,
+  link: AnyLink,
+  options: LinkRenderOptions = {},
+) {
+  return `${formatKey(name, options.keyOptions)}: ${renderLinkValue(
+    link,
+    options,
+  )}`;
+}
+
+export function renderLinksObject(
+  links: LinksDef<any>,
+  options: LinksObjectOptions = {},
+) {
+  const indent = options.indent ?? DEFAULT_INDENT;
+  const entries = sortedEntries(links).map(([name, link]) =>
+    renderLinkProperty(name, link, {
+      indent,
+      keyOptions: options.keyOptions,
+      stringQuote: options.stringQuote,
+      style: options.style,
+    }),
+  );
+  if (!entries.length && !options.newlineForEmpty) {
+    return '{}';
+  }
+  const inner = entries.length ? indentLines(entries.join(',\n'), indent) : '';
+  return `{\n${inner}\n}`;
+}
+
+function isValidIdentifier(name: string) {
+  return /^[A-Za-z_$][A-Za-z0-9_$]*$/.test(name);
+}