diff --git a/src/plus/webviews/graph/graphWebview.ts b/src/plus/webviews/graph/graphWebview.ts
index b47f16a..3c1662e 100644
--- a/src/plus/webviews/graph/graphWebview.ts
+++ b/src/plus/webviews/graph/graphWebview.ts
@@ -65,6 +65,7 @@ import type {
 	GetMoreCommitsParams,
 	GraphColumnConfig,
 	GraphColumnName,
+	GraphColumnsSettings,
 	GraphComponentConfig,
 	GraphRepository,
 	SearchCommitsParams,
@@ -102,6 +103,15 @@ export interface GraphSelectionChangeEvent {
 	readonly selection: GitCommit[];
 }
 
+const defaultGraphColumnsSettings: GraphColumnsSettings = {
+	ref: { width: 150, isHidden: false },
+	graph: { width: 150, isHidden: false },
+	message: { width: 300, isHidden: false },
+	author: { width: 130, isHidden: false },
+	datetime: { width: 130, isHidden: false },
+	sha: { width: 130, isHidden: false },
+};
+
 export class GraphWebview extends WebviewBase<State> {
 	private _onDidChangeSelection = new EventEmitter<GraphSelectionChangeEvent>();
 	get onDidChangeSelection(): Event<GraphSelectionChangeEvent> {
@@ -692,7 +702,7 @@ export class GraphWebview extends WebviewBase<State> {
 
 		const columns = this.getColumns();
 		return this.notify(DidChangeColumnsNotificationType, {
-			columns: columns,
+			columns: this.getColumnSettings(columns),
 			context: this.getColumnHeaderContext(columns),
 		});
 	}
@@ -831,6 +841,26 @@ export class GraphWebview extends WebviewBase<State> {
 		return this.container.storage.getWorkspace('graph:columns');
 	}
 
+	private getColumnSettings(
+		columns: Record<GraphColumnName, GraphColumnConfig> | undefined,
+	): GraphColumnsSettings | undefined {
+		if (columns == null) return undefined;
+
+		const columnsSettings: GraphColumnsSettings = {
+			...defaultGraphColumnsSettings,
+		};
+		if (columns != null) {
+			for (const [column, columnCfg] of Object.entries(columns) as [GraphColumnName, GraphColumnConfig][]) {
+				columnsSettings[column] = {
+					...defaultGraphColumnsSettings[column],
+					...columnCfg,
+				};
+			}
+		}
+
+		return columnsSettings;
+	}
+
 	private getColumnHeaderContext(columns: Record<GraphColumnName, GraphColumnConfig> | undefined): string {
 		const hidden: string[] = [];
 		if (columns != null) {
@@ -945,7 +975,7 @@ export class GraphWebview extends WebviewBase<State> {
 							hasMore: data.paging?.hasMore ?? false,
 					  }
 					: undefined,
-			columns: columns,
+			columns: this.getColumnSettings(columns),
 			config: this.getComponentConfig(),
 			context: {
 				header: this.getColumnHeaderContext(columns),
diff --git a/src/plus/webviews/graph/protocol.ts b/src/plus/webviews/graph/protocol.ts
index ae13c4b..72cdfde 100644
--- a/src/plus/webviews/graph/protocol.ts
+++ b/src/plus/webviews/graph/protocol.ts
@@ -26,7 +26,7 @@ export interface State {
 	loading?: boolean;
 	rows?: GraphRow[];
 	paging?: GraphPaging;
-	columns?: Record<GraphColumnName, GraphColumnConfig>;
+	columns?: GraphColumnsSettings;
 	config?: GraphComponentConfig;
 	context?: GraphContexts;
 	nonce?: string;
@@ -168,7 +168,7 @@ export const DidChangeAvatarsNotificationType = new IpcNotificationType<DidChang
 );
 
 export interface DidChangeColumnsParams {
-	columns: Record<GraphColumnName, GraphColumnConfig> | undefined;
+	columns: GraphColumnsSettings | undefined;
 	context?: string;
 }
 export const DidChangeColumnsNotificationType = new IpcNotificationType<DidChangeColumnsParams>(
diff --git a/src/webviews/apps/plus/graph/GraphWrapper.tsx b/src/webviews/apps/plus/graph/GraphWrapper.tsx
index ef4150a..d8d9794 100644
--- a/src/webviews/apps/plus/graph/GraphWrapper.tsx
+++ b/src/webviews/apps/plus/graph/GraphWrapper.tsx
@@ -18,7 +18,6 @@ import type {
 	DismissBannerParams,
 	GraphColumnConfig,
 	GraphColumnName,
-	GraphColumnsSettings,
 	GraphComponentConfig,
 	GraphRepository,
 	State,
@@ -74,30 +73,6 @@ const getStyleProps = (
 	};
 };
 
-const defaultGraphColumnsSettings: GraphColumnsSettings = {
-	ref: { width: 150, isHidden: false },
-	graph: { width: 150, isHidden: false },
-	message: { width: 300, isHidden: false },
-	author: { width: 130, isHidden: false },
-	datetime: { width: 130, isHidden: false },
-	sha: { width: 130, isHidden: false },
-};
-
-const getGraphColumns = (columns?: Record<GraphColumnName, GraphColumnConfig> | undefined): GraphColumnsSettings => {
-	const columnsSettings: GraphColumnsSettings = {
-		...defaultGraphColumnsSettings,
-	};
-	if (columns != null) {
-		for (const [column, columnCfg] of Object.entries(columns) as [GraphColumnName, GraphColumnConfig][]) {
-			columnsSettings[column] = {
-				...defaultGraphColumnsSettings[column],
-				...columnCfg,
-			};
-		}
-	}
-	return columnsSettings;
-};
-
 const getGraphDateFormatter = (config?: GraphComponentConfig): OnFormatCommitDateTime => {
 	return (commitDateTime: number) => formatCommitDateTime(commitDateTime, config?.dateStyle, config?.dateFormat);
 };
@@ -206,7 +181,7 @@ export function GraphWrapper({
 	const [graphSelectedRows, setSelectedRows] = useState(selectedRows);
 	const [graphConfig, setGraphConfig] = useState(config);
 	// const [graphDateFormatter, setGraphDateFormatter] = useState(getGraphDateFormatter(config));
-	const [graphColumns, setGraphColumns] = useState(getGraphColumns(columns));
+	const [graphColumns, setGraphColumns] = useState(columns);
 	const [graphContext, setGraphContext] = useState(context);
 	const [pagingState, setPagingState] = useState(paging);
 	const [isLoading, setIsLoading] = useState(loading);
@@ -409,7 +384,7 @@ export function GraphWrapper({
 		}
 		setGraphConfig(state.config);
 		// setGraphDateFormatter(getGraphDateFormatter(config));
-		setGraphColumns(getGraphColumns(state.columns));
+		setGraphColumns(state.columns);
 		setGraphContext(state.context);
 		setPagingState(state.paging);
 		setStyleProps(getStyleProps(state.mixedColumnColors));
diff --git a/src/webviews/apps/plus/graph/graph.tsx b/src/webviews/apps/plus/graph/graph.tsx
index 1ef2372..75acecc 100644
--- a/src/webviews/apps/plus/graph/graph.tsx
+++ b/src/webviews/apps/plus/graph/graph.tsx
@@ -120,13 +120,18 @@ export class GraphApp extends App<State> {
 
 			case DidChangeColumnsNotificationType.method:
 				onIpc(DidChangeColumnsNotificationType, msg, params => {
-					this.setState({
-						...this.state,
-						...(params.columns != null ? { columns: params.columns } : undefined),
-						...(params.context != null
-							? { context: { ...this.state.context, header: params.context } }
-							: undefined),
-					});
+					const newState = { ...this.state, columns: params.columns };
+					if (params.context != null) {
+						if (newState.context == null) {
+							newState.context = { header: params.context };
+						} else {
+							newState.context.header = params.context;
+						}
+					} else if (newState.context?.header != null) {
+						newState.context.header = undefined;
+					}
+
+					this.setState(newState);
 					this.refresh(this.state);
 				});
 				break;