diff --git a/src/index.ts b/src/index.ts
index ed6aeb7..ed0fd40 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -115,23 +115,26 @@ function amsmathBlock(
   let { endpos } = outcome
   endpos += begin
 
-  let line = startLine
-  while (line < endLine) {
-    if (endpos >= state.bMarks[line] && endpos <= state.eMarks[line]) {
+  // Find the final line of the amsmath block
+  let nextLine = startLine
+  for (; nextLine < endLine; nextLine++) {
+    if (endpos >= state.bMarks[nextLine] && endpos <= state.eMarks[nextLine]) {
       // line for end of block math found ...
       // eslint-disable-next-line no-param-reassign
-      state.line = line + 1
       break
     }
-    line += 1
   }
+  // Skip over the final amsmath line
+  nextLine += 1;
+  // Write back to the state
+  state.line = nextLine;
 
   if (!silent) {
     const token = state.push("amsmath", "math", 0)
     token.block = true
     token.content = state.src.slice(begin, endpos)
     token.meta = { environment, numbered }
-    token.map = [startLine, line]
+    token.map = [startLine, nextLine]
   }
 
   return true
diff --git a/tests/fixtures.spec.ts b/tests/fixtures.spec.ts
index 64ac1c9..4fa19a8 100644
--- a/tests/fixtures.spec.ts
+++ b/tests/fixtures.spec.ts
@@ -1,4 +1,4 @@
-import { describe, expect, it } from "vitest"
+import { describe, expect, test, it } from "vitest"
 import fs from "node:fs"
 import katex from "katex"
 import MarkdownIt from "markdown-it"
@@ -37,3 +37,56 @@ describe("Parses basic", () => {
     it(name, () => expect(rendered.trim()).toEqual(expected.trim()))
   })
 })
+
+test("Produces correct tokens", () => {
+  const mdit = MarkdownIt().use(amsmathPlugin)
+  const tokens = mdit.parse(`
+\\begin{equation}
+y = m * x + c \\
+z = 2
+\\end{equation}
+Hello world!
+`)
+  expect(tokens).toMatchObject([
+    {
+      type: "amsmath",
+      map: [1, 5],
+      nesting: 0,
+      level: 0,
+      content: "\\begin{equation}\ny = m * x + c \\\nz = 2\n\\end{equation}",
+      meta: { environment: "equation", numbered: "" },
+      block: true,
+      hidden: false
+    },
+    {
+      type: "paragraph_open",
+      map: [5, 6],
+      nesting: 1,
+      level: 0,
+      content: "",
+      meta: null,
+      block: true,
+      hidden: false
+    },
+    {
+      type: "inline",
+      map: [5, 6],
+      nesting: 0,
+      level: 1,
+      content: "Hello world!",
+      meta: null,
+      block: true,
+      hidden: false
+    },
+    {
+      type: "paragraph_close",
+      map: null,
+      nesting: -1,
+      level: 0,
+      content: "",
+      meta: null,
+      block: true,
+      hidden: false
+    }
+  ])
+})