diff --git a/src/main/java/de/prob2/jupyter/commands/DotCommand.java b/src/main/java/de/prob2/jupyter/commands/DotCommand.java
index ac43f778a61a35f9df937d440290c52ef5d1656d..53ee6926423c26dbe6d0f420b157e8edf52a93d5 100644
--- a/src/main/java/de/prob2/jupyter/commands/DotCommand.java
+++ b/src/main/java/de/prob2/jupyter/commands/DotCommand.java
@@ -99,7 +99,13 @@ public final class DotCommand implements Command {
 			throw new UncheckedIOException("Failed to create temp file", e);
 		}
 		final GetSvgForVisualizationCommand cmd2 = new GetSvgForVisualizationCommand(trace.getCurrentState(), item, outPath.toFile(), args);
-		trace.getStateSpace().execute(cmd2);
+		// Provide source code (if any) to error highlighter
+		final Runnable execute = () -> trace.getStateSpace().execute(cmd2);
+		if (split.size() > 1) {
+			CommandUtils.withSourceCode(split.get(1), execute);
+		} else {
+			execute.run();
+		}
 		final String svg;
 		try (final Stream<String> lines = Files.lines(outPath)) {
 			svg = lines.collect(Collectors.joining("\n"));