diff --git a/flask_app/app.py b/flask_app/app.py
index fa392df..a0838a0 100644
--- a/flask_app/app.py
+++ b/flask_app/app.py
@@ -403,8 +403,14 @@ def run_demo(module_name: str) -> str:
 
     buf = io.StringIO()
     with contextlib.redirect_stdout(buf):
-        demo_fn()
-    return buf.getvalue()
+        result = demo_fn()
+
+    # If demo() printed to stdout, use that; otherwise use return value if present
+    output = buf.getvalue()
+    if not output and result is not None:
+        output = str(result)
+
+    return output
 
 
 @app.get("/demo")
diff --git a/flask_app/templates/demo.html b/flask_app/templates/demo.html
index 3ae5cee..c4a339b 100644
--- a/flask_app/templates/demo.html
+++ b/flask_app/templates/demo.html
@@ -26,10 +26,14 @@ <h3>Algorithm Notes</h3>
       </section>
     {% endif %}
 
-    {% if output %}
+    {% if output is not none %}
       <section class="panel">
         <h3>Output</h3>
-        <pre class="output"><code>{{ output }}</code></pre>
+        {% if output %}
+          <pre class="output"><code>{{ output }}</code></pre>
+        {% else %}
+          <p class="muted">Demo executed successfully with no output.</p>
+        {% endif %}
       </section>
     {% endif %}