PolicyEngine · MaxGhenis · May 10, 2026 · May 10, 2026
diff --git a/changelog.d/shell-followups.fixed.md b/changelog.d/shell-followups.fixed.md
@@ -0,0 +1 @@
+Preserve PolicyEngine shell client boundaries and fix default shell navigation edge cases.
diff --git a/src/layout/PolicyEngineHeader.tsx b/src/layout/PolicyEngineHeader.tsx
@@ -3,6 +3,7 @@
 import { Header, type CountryConfig, type HeaderProps, type NavItemConfig } from './header';
 import {
   DEFAULT_POLICYENGINE_BASE_URL,
+  getPolicyEngineCountrySwitchUrl,
   getPolicyEngineCountryUrl,
   getPolicyEngineNavItems,
   policyEngineCountries,
@@ -35,7 +36,10 @@ export function PolicyEngineHeader({
     onCountryChange ??
     ((countryId: string) => {
       if (typeof window !== 'undefined') {
-        window.location.href = getPolicyEngineCountryUrl(countryId, baseUrl);
+        window.location.href = getPolicyEngineCountrySwitchUrl(countryId, {
+          baseUrl,
+          countries,
+        });
       }
     });
 

diff --git a/src/layout/PolicyEngineSiteConfig.ts b/src/layout/PolicyEngineSiteConfig.ts
@@ -21,6 +21,35 @@ export function getPolicyEngineCountryUrl(
   return `${normalizeBaseUrl(baseUrl)}/${country}`;
 }
 
+export interface PolicyEngineCountrySwitchUrlOptions {
+  baseUrl?: string;
+  countries?: CountryConfig[];
+  pathname?: string;
+  search?: string;
+  hash?: string;
+}
+
+export function getPolicyEngineCountrySwitchUrl(
+  country: PolicyEngineCountryId,
+  {
+    baseUrl = DEFAULT_POLICYENGINE_BASE_URL,
+    countries = policyEngineCountries,
+    pathname = typeof window !== 'undefined' ? window.location.pathname : '',
+    search = typeof window !== 'undefined' ? window.location.search : '',
+    hash = typeof window !== 'undefined' ? window.location.hash : '',
+  }: PolicyEngineCountrySwitchUrlOptions = {},
+) {
+  const countryIds = new Set(countries.map(({ id }) => id));
+  const segments = pathname.split('/');
+
+  if (countryIds.has(segments[1])) {
+    segments[1] = country;
+    return `${normalizeBaseUrl(baseUrl)}${segments.join('/')}${search}${hash}`;
+  }
+
+  return getPolicyEngineCountryUrl(country, baseUrl);
+}
+
 export function getPolicyEngineUrl(
   country: PolicyEngineCountryId = 'us',
   path = '',
@@ -34,7 +63,7 @@ export function getPolicyEngineNavItems(
   country: PolicyEngineCountryId = 'us',
   baseUrl = DEFAULT_POLICYENGINE_BASE_URL,
 ): NavItemConfig[] {
-  return [
+  const navItems: NavItemConfig[] = [
     {
       label: 'Research',
       href: getPolicyEngineUrl(country, 'research', baseUrl),
@@ -47,10 +76,6 @@ export function getPolicyEngineNavItems(
       label: 'API',
       href: getPolicyEngineUrl(country, 'api', baseUrl),
     },
-    {
-      label: 'Python',
-      href: getPolicyEngineUrl(country, 'python', baseUrl),
-    },
     {
       label: 'About',
       href: getPolicyEngineUrl(country, 'team', baseUrl),
@@ -77,6 +102,15 @@ export function getPolicyEngineNavItems(
       href: getPolicyEngineUrl(country, 'donate', baseUrl),
     },
   ];
+
+  if (country === 'us') {
+    navItems.splice(3, 0, {
+      label: 'Python',
+      href: getPolicyEngineUrl(country, 'python', baseUrl),
+    });
+  }
+
+  return navItems;
 }
 
 export function getPolicyEngineFooterLinks(

diff --git a/src/layout/header/HeaderNavItem.tsx b/src/layout/header/HeaderNavItem.tsx
@@ -1,4 +1,4 @@
-import { useCallback, useEffect, useRef, useState } from 'react';
+import { useEffect, useRef, useState } from 'react';
 import { ChevronDown } from 'lucide-react';
 import type { NavItemConfig } from './Header';
 
@@ -29,6 +29,27 @@ const hoverHandlers = {
   },
 };
 
+const dropdownItemStyle: React.CSSProperties = {
+  display: 'flex',
+  alignItems: 'center',
+  width: '100%',
+  textAlign: 'left',
+  padding: '11px 16px',
+  borderRadius: '10px',
+  border: 'none',
+  background: 'transparent',
+  cursor: 'pointer',
+  fontSize: '14px',
+  fontFamily: 'var(--font-sans)',
+  fontWeight: 600,
+  color: 'var(--color-teal-800)',
+  transition:
+    'background-color 0.12s ease, color 0.12s ease, opacity 0.3s ease',
+  lineHeight: '1.3',
+  letterSpacing: '-0.01em',
+  textDecoration: 'none',
+};
+
 /**
  * Apple-style dropdown panel with smooth height reveal and content fade.
  */
@@ -60,20 +81,6 @@ function AppleDropdown({
     }
   }, [open]);
 
-  const handleSelect = useCallback(
-    (item: { label: string; href: string }) => {
-      onClose();
-      if (onNavigate) {
-        onNavigate(item.href);
-      } else if (linkComponent) {
-        // linkComponent handles its own navigation
-      } else {
-        window.location.href = item.href;
-      }
-    },
-    [onClose, onNavigate, linkComponent],
-  );
-
   if (!open && contentHeight === 0) {
     return null;
   }
@@ -118,57 +125,65 @@ function AppleDropdown({
         }}
       >
         <div ref={contentRef} style={{ padding: '8px' }}>
-          {items.map((item, i) => (
-            <button
-              key={item.label}
-              type="button"
-              onClick={() => handleSelect(item)}
-              style={{
-                display: 'flex',
-                alignItems: 'center',
-                width: '100%',
-                textAlign: 'left',
-                padding: '11px 16px',
-                borderRadius: '10px',
-                border: 'none',
-                background: 'transparent',
-                cursor: 'pointer',
-                fontSize: '14px',
-                fontFamily: 'var(--font-sans)',
-                fontWeight: 600,
-                color: 'var(--color-teal-800)',
-                transition:
-                  'background-color 0.12s ease, color 0.12s ease, opacity 0.3s ease',
-                transitionDelay: visible ? `${i * 50}ms` : '0ms',
-                opacity: visible ? 1 : 0,
-                lineHeight: '1.3',
-                letterSpacing: '-0.01em',
-              }}
-              onMouseEnter={(e) => {
-                e.currentTarget.style.backgroundColor =
-                  'var(--color-teal-500)';
-                e.currentTarget.style.color = 'var(--text-inverse)';
-              }}
-              onMouseLeave={(e) => {
-                e.currentTarget.style.backgroundColor = 'transparent';
-                e.currentTarget.style.color = 'var(--color-teal-800)';
-              }}
-            >
-              <span>{item.label}</span>
-              {item.description && (
-                <span
-                  style={{
-                    marginLeft: '8px',
-                    fontSize: '12px',
-                    opacity: 0.6,
-                    fontWeight: 400,
-                  }}
+          {items.map((item, i) => {
+            const style = {
+              ...dropdownItemStyle,
+              transitionDelay: visible ? `${i * 50}ms` : '0ms',
+              opacity: visible ? 1 : 0,
+            };
+            const content = (
+              <>
+                <span>{item.label}</span>
+                {item.description && (
+                  <span
+                    style={{
+                      marginLeft: '8px',
+                      fontSize: '12px',
+                      opacity: 0.6,
+                      fontWeight: 400,
+                    }}
+                  >
+                    {item.description}
+                  </span>
+                )}
+              </>
+            );
+            const handleClick = (e: React.MouseEvent<HTMLElement>) => {
+              onClose();
+              if (onNavigate) {
+                e.preventDefault();
+                onNavigate(item.href);
+              }
+            };
+
+            if (linkComponent) {
+              const LinkComp = linkComponent;
+              return (
+                <LinkComp
+                  key={item.label}
+                  href={item.href}
+                  to={item.href}
+                  onClick={handleClick}
+                  style={style}
+                  {...hoverHandlers}
                 >
-                  {item.description}
-                </span>
-              )}
-            </button>
-          ))}
+                  {content}
+                </LinkComp>
+              );
+            }
+
+            return (
+              <a
+                key={item.label}
+                href={item.href}
+                onClick={handleClick}
+                style={style}
+                {...hoverHandlers}
+              >
+                {content}
+              </a>
+            );
+          })}
         </div>
       </div>
     </>

diff --git a/src/layout/index.ts b/src/layout/index.ts
@@ -24,6 +24,7 @@ export {
 } from './PolicyEngineShell';
 export {
   DEFAULT_POLICYENGINE_BASE_URL,
+  getPolicyEngineCountrySwitchUrl,
   getPolicyEngineCountryUrl,
   getPolicyEngineFooterLinks,
   getPolicyEngineNavItems,

diff --git a/tests/consumer-types/client-directives.test.ts b/tests/consumer-types/client-directives.test.ts
@@ -0,0 +1,23 @@
+import * as fs from 'node:fs';
+import * as path from 'node:path';
+import { describe, expect, it } from 'vitest';
+
+const ROOT = path.resolve(__dirname, '..', '..');
+const DIST_INDEX = path.join(ROOT, 'dist', 'index.js');
+const DIST_LAYOUT = path.join(ROOT, 'dist', 'layout.js');
+
+const distAvailable =
+  fs.existsSync(DIST_INDEX) && fs.existsSync(DIST_LAYOUT);
+
+describe('client directives in built artifacts', () => {
+  if (!distAvailable) {
+    it.skip('(skipped — run `bun run build` first to inspect dist/ artifacts)', () => {});
+    return;
+  }
+
+  it('marks shell entrypoints as client modules for Next App Router consumers', () => {
+    for (const file of [DIST_INDEX, DIST_LAYOUT]) {
+      expect(fs.readFileSync(file, 'utf8').startsWith('"use client";')).toBe(true);
+    }
+  });
+});
diff --git a/tests/layout/Layout.test.tsx b/tests/layout/Layout.test.tsx
@@ -1,12 +1,14 @@
 import { describe, it, expect } from 'vitest';
 import { render, screen } from '@testing-library/react';
+import userEvent from '@testing-library/user-event';
 import { DashboardShell } from '../../src/layout/DashboardShell';
 import { SidebarLayout } from '../../src/layout/SidebarLayout';
 import { SingleColumnLayout } from '../../src/layout/SingleColumnLayout';
 import { Header } from '../../src/layout/header';
 import { InputPanel } from '../../src/layout/InputPanel';
 import {
   getPolicyEngineNavItems,
+  getPolicyEngineCountrySwitchUrl,
   PolicyEngineFooter,
   PolicyEngineHeader,
   PolicyEngineShell,
@@ -58,6 +60,19 @@ describe('PolicyEngine shell components', () => {
       label: 'API',
       href: 'https://policyengine.org/uk/api',
     });
+    expect(navItems.some((item) => item.label === 'Python')).toBe(false);
+  });
+
+  it('preserves the current route when building country switch URLs', () => {
+    expect(
+      getPolicyEngineCountrySwitchUrl('uk', {
+        pathname: '/us/california-wealth-tax/results',
+        search: '?household=1',
+        hash: '#chart',
+      }),
+    ).toBe(
+      'https://policyengine.org/uk/california-wealth-tax/results?household=1#chart',
+    );
   });
 
   it('renders the canonical PolicyEngine header', () => {
@@ -81,6 +96,26 @@ describe('PolicyEngine shell components', () => {
     expect(screen.getByText('Tool content')).toBeInTheDocument();
     expect(screen.getAllByText('Donate').length).toBeGreaterThan(0);
   });
+
+  it('renders dropdown links with a custom link component', async () => {
+    const Link = ({
+      children,
+      href,
+      ...props
+    }: React.AnchorHTMLAttributes<HTMLAnchorElement> & { href: string }) => (
+      <a href={href} {...props}>
+        {children}
+      </a>
+    );
+
+    render(<PolicyEngineHeader country="us" linkComponent={Link} />);
+
+    await userEvent.click(screen.getByRole('button', { name: /about/i }));
+    expect(screen.getByRole('link', { name: /Team/i })).toHaveAttribute(
+      'href',
+      'https://policyengine.org/us/team',
+    );
+  });
 });
 
 describe('InputPanel', () => {

diff --git a/vite.config.ts b/vite.config.ts
@@ -2,9 +2,34 @@ import { resolve } from 'path';
 import tailwindcss from '@tailwindcss/vite';
 import react from '@vitejs/plugin-react';
 import { defineConfig } from 'vite';
+import type { Plugin } from 'vite';
+
+function preserveClientDirectives(): Plugin {
+  return {
+    name: 'preserve-client-directives',
+    generateBundle(_options, bundle) {
+      for (const chunk of Object.values(bundle)) {
+        if (chunk.type !== 'chunk') continue;
+
+        const isClientEntry =
+          chunk.facadeModuleId?.endsWith('/src/index.ts') ||
+          chunk.facadeModuleId?.endsWith('/src/layout/index.ts');
+        const hasClientLayoutModule = chunk.moduleIds.some(
+          (id) =>
+            id.includes('/src/layout/PolicyEngine') ||
+            id.includes('/src/layout/header/'),
+        );
+
+        if ((isClientEntry || hasClientLayoutModule) && !chunk.code.startsWith('"use client";')) {
+          chunk.code = `"use client";\n${chunk.code}`;
+        }
+      }
+    },
+  };
+}
 
 export default defineConfig({
-  plugins: [tailwindcss(), react()],
+  plugins: [tailwindcss(), react(), preserveClientDirectives()],
   resolve: {
     alias: {
       '@': resolve(__dirname, 'src'),
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1 @@
		Preserve PolicyEngine shell client boundaries and fix default shell navigation edge cases.