Turn AI-download chrome://flags into per-item options with select-all and revert

README.md

@@ -67,7 +67,9 @@ It can also **block on-device AI model downloads** (Gemini Nano), which is the d
   - Linux: `/etc/opt/chrome/policies/managed/go-chrome-ai.json` (needs sudo)
   - Windows: `HKLM\Software\Policies\Google\Chrome` REG_DWORD
 
-Because the third change is an Enterprise policy, Chrome will display the "managed by your organization" banner afterwards. Pass `-disable-ai-download=false` (CLI) or untick the option (GUI) if you do not want that.
+Because the third change is an Enterprise policy, Chrome will display the "managed by your organization" banner afterwards. All three are independently selectable options — the two `chrome://flags` entries (CLI: `-disable-flag`, GUI: per-flag checkboxes) and the `chrome://policy` write (CLI: `-disable-ai-policy`, GUI: policy checkbox) — or all three at once via the "select all" option (CLI: `-disable-ai-download`, GUI: "Select all" checkbox, both default on). Turn off `-disable-ai-download` (CLI) or "Select all" (GUI) to pick and choose, e.g. to disable the flags without triggering the Enterprise-policy banner, or vice versa.
+
+Every run fully syncs Chrome to the current selection, in both directions: a selected item is applied (flag forced to Disabled / policy written), and a **deselected item is actively reverted** — its `chrome://flags` override is removed (back to Chrome's default) and the `chrome://policy` entry is deleted from the OS managed-policy store, if either was previously set by this tool. Unchecking an option is not a no-op; it undoes that option's effect on the next run.
 
 ## Screenshot
 
@@ -88,7 +90,9 @@ Flags:
 
 - `-dry-run`: show changes without writing files or killing Chrome
 - `-no-restart`: patch but do not restart Chrome
-- `-disable-ai-download` (default `true`): block on-device AI model downloads by disabling the relevant `chrome://flags` entries and writing `GenAILocalFoundationalModelSettings=1` to the OS managed-policy store. Use `-disable-ai-download=false` to skip.
+- `-disable-ai-download` (default `true`): "select all" — block on-device AI model downloads by disabling every known `chrome://flags` entry and writing `GenAILocalFoundationalModelSettings=1` to the OS managed-policy store. Use `-disable-ai-download=false` to pick individually with `-disable-flag` and/or `-disable-ai-policy` instead.
+- `-disable-flag <name>` (repeatable): disable one specific `chrome://flags` entry by name (e.g. `optimization-guide-on-device-model`, `prompt-api-for-gemini-nano`). Only takes effect when `-disable-ai-download=false`.
+- `-disable-ai-policy` (default `true`): write the `GenAILocalFoundationalModelSettings` Enterprise policy, independent of which flags are selected. Only takes effect when `-disable-ai-download=false`.
 
 ## Run GUI
 
@@ -105,7 +109,7 @@ The GUI includes:
 - one-click patch flow
 - progress bar
 - real-time logs
-- a "Disable on-device AI model download" toggle that previews the exact `chrome://flags` and `chrome://policy` changes before you press Run
+- a checkbox per `chrome://flags` entry, a checkbox for the `chrome://policy` write, and a "Select all" master switch over both (checking it selects and locks everything; uncheck it to pick items individually), with a live preview of the exact changes before you press Run
 
 ## Build From Source
 

cmd/go-chrome-ai/main.go

@@ -33,6 +33,17 @@ func printUsage() {
 	fmt.Println("  go-chrome-ai cli [flags]")
 	fmt.Println("")
 	fmt.Println("Flags:")
-	fmt.Println("  -dry-run     Show what would change without modifying files")
-	fmt.Println("  -no-restart  Do not restart Chrome after patching")
+	fmt.Println("  -dry-run              Show what would change without modifying files")
+	fmt.Println("  -no-restart           Do not restart Chrome after patching")
+	fmt.Println("  -disable-ai-download  Block on-device Gemini Nano download by disabling all")
+	fmt.Println("                        known chrome://flags entries plus the OS policy (default true)")
+	fmt.Println("  -disable-flag         Individual chrome://flags entry to disable (repeatable);")
+	fmt.Println("                        only takes effect when -disable-ai-download=false")
+	fmt.Println("  -disable-ai-policy    Write the GenAILocalFoundationalModelSettings Enterprise")
+	fmt.Println("                        policy (chrome://policy), independent of -disable-flag;")
+	fmt.Println("                        only takes effect when -disable-ai-download=false (default true)")
+	fmt.Println("")
+	fmt.Println("Every run syncs to the current selection: a deselected flag or policy is")
+	fmt.Println("actively reverted (chrome://flags reset to default, chrome://policy entry")
+	fmt.Println("removed) if this tool previously set it.")
 }

internal/app/cli.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"io"
 	"os"
+	"strings"
 
 	"github.com/itamaker/go-chrome-ai/internal/chrome"
 	"github.com/itamaker/go-chrome-ai/internal/meta"
@@ -21,10 +22,26 @@ func RunCLI(args []string, stderr io.Writer) int {
 	fs := flag.NewFlagSet("go-chrome-ai", flag.ContinueOnError)
 	fs.SetOutput(stderr)
 
+	availableFlagNames := chrome.AllAIDownloadFlagNames()
+
 	dryRun := fs.Bool("dry-run", false, "Show what would change without modifying files")
 	noRestart := fs.Bool("no-restart", false, "Do not restart Chrome after patching")
-	disableAI := fs.Bool("disable-ai-download", true,
-		"Block on-device Gemini Nano download (use -disable-ai-download=false to skip)")
+	selectAllAIFlags := fs.Bool("disable-ai-download", true,
+		"Block on-device Gemini Nano download by disabling all known chrome://flags entries "+
+			"("+strings.Join(availableFlagNames, ", ")+") plus writing the OS Enterprise policy. "+
+			"Use -disable-ai-download=false to pick individual flags with -disable-flag and/or "+
+			"the policy with -disable-ai-policy instead.")
+	var selectedFlagNames []string
+	fs.Func("disable-flag",
+		"Individual chrome://flags entry to disable (repeatable); only takes effect when -disable-ai-download=false",
+		func(v string) error {
+			selectedFlagNames = append(selectedFlagNames, v)
+			return nil
+		})
+	applyPolicy := fs.Bool("disable-ai-policy", true,
+		"Write the "+chrome.GenAIPolicyName+" Enterprise policy (chrome://policy) that also blocks "+
+			"on-device AI downloads; causes Chrome to show the \"managed by your organization\" banner. "+
+			"Independent of which chrome://flags are disabled. Only takes effect when -disable-ai-download=false.")
 
 	if err := fs.Parse(args); err != nil {
 		if errors.Is(err, flag.ErrHelp) {
@@ -33,20 +50,45 @@ func RunCLI(args []string, stderr io.Writer) int {
 		return 2
 	}
 
-	if *disableAI {
-		fmt.Println("Disable AI model download - will apply:")
-		fmt.Println("  Local flag overrides (chrome://flags):")
-		for _, action := range chrome.DisableAIDownloadActions() {
-			if action.EnterprisePolicy {
-				continue
+	var aiDownloadFlags []string
+	var aiDownloadPolicy bool
+	if *selectAllAIFlags {
+		aiDownloadFlags = availableFlagNames
+		aiDownloadPolicy = true
+	} else {
+		known := make(map[string]bool, len(availableFlagNames))
+		for _, name := range availableFlagNames {
+			known[name] = true
+		}
+		for _, name := range selectedFlagNames {
+			if !known[name] {
+				fmt.Fprintf(stderr, "Error: unknown -disable-flag %q (known: %s)\n",
+					name, strings.Join(availableFlagNames, ", "))
+				return 2
 			}
+		}
+		aiDownloadFlags = selectedFlagNames
+		aiDownloadPolicy = *applyPolicy
+	}
+
+	actions := chrome.DisableAIDownloadActions(aiDownloadFlags, aiDownloadPolicy)
+	applyFlags, revertFlags, policyActions := chrome.GroupDisableAIDownloadActions(actions)
+	fmt.Println("Chrome AI-download configuration:")
+	if len(applyFlags) > 0 {
+		fmt.Println("  Local flag overrides (chrome://flags):")
+		for _, action := range applyFlags {
+			fmt.Println("    - " + action.Label)
+		}
+	}
+	if len(revertFlags) > 0 {
+		fmt.Println("  Local flag resets (chrome://flags):")
+		for _, action := range revertFlags {
 			fmt.Println("    - " + action.Label)
 		}
+	}
+	if len(policyActions) > 0 {
 		fmt.Println("  Enterprise policy (chrome://policy):")
-		for _, action := range chrome.DisableAIDownloadActions() {
-			if !action.EnterprisePolicy {
-				continue
-			}
+		for _, action := range policyActions {
 			fmt.Println("    - " + action.Label)
 			if action.Detail != "" {
 				fmt.Println("      " + action.Detail)
@@ -58,9 +100,10 @@ func RunCLI(args []string, stderr io.Writer) int {
 	}
 
 	summary, err := chrome.Run(chrome.Options{
-		DryRun:                 *dryRun,
-		NoRestart:              *noRestart,
-		DisableAIModelDownload: *disableAI,
+		DryRun:           *dryRun,
+		NoRestart:        *noRestart,
+		AIDownloadFlags:  aiDownloadFlags,
+		AIDownloadPolicy: aiDownloadPolicy,
 	}, chrome.Callbacks{
 		Log: func(message string) {
 			fmt.Println(message)

internal/chrome/flags.go

@@ -7,47 +7,108 @@ import "strings"
 // `<flag-name>@<choice>` where 0=Default, 1=Enabled, 2=Disabled.
 const flagDisabledSuffix = "@2"
 
-// AIDownloadFlagNames are the chrome://flags entries this tool forces to
+// AIDownloadFlag describes one chrome://flags entry this tool can force to
 // "Disabled" so Chrome does not download Gemini Nano / on-device models.
-var AIDownloadFlagNames = []string{
-	"optimization-guide-on-device-model",
-	"prompt-api-for-gemini-nano",
+type AIDownloadFlag struct {
+	Name string // chrome://flags entry name
 }
 
-// DisableAIDownloadAction describes one transform applied by the
+// AvailableAIDownloadFlags lists every chrome://flags entry this tool knows
+// how to disable. Callers (CLI/GUI) present each one as an individually
+// selectable option, plus a "select all" convenience over this same list.
+var AvailableAIDownloadFlags = []AIDownloadFlag{
+	{Name: "optimization-guide-on-device-model"},
+	{Name: "prompt-api-for-gemini-nano"},
+}
+
+// AllAIDownloadFlagNames returns the names of every available AI-download
+// flag, i.e. the set selected by a "select all" option.
+func AllAIDownloadFlagNames() []string {
+	names := make([]string, len(AvailableAIDownloadFlags))
+	for i, f := range AvailableAIDownloadFlags {
+		names[i] = f.Name
+	}
+	return names
+}
+
+// DisableAIDownloadAction describes one transform previewed/applied by the
 // "disable AI model download" feature.
 type DisableAIDownloadAction struct {
 	Label            string // human-readable label, e.g. "chrome://flags/#foo -> Disabled"
 	Detail           string // optional second line (e.g. policy storage location)
-	EnterprisePolicy bool   // true if the action writes a managed Chrome policy
+	EnterprisePolicy bool   // true if the action reads/writes the managed Chrome policy
 	PolicyNote       string // extra warning shown when EnterprisePolicy is true
+	Revert           bool   // true if this action undoes a previous change instead of applying one
 }
 
-// DisableAIDownloadActions returns the ordered list of changes applied when
-// the user enables "disable AI model download". The last entry is an
-// Enterprise policy write that causes Chrome to display the
-// "managed by your organization" banner.
-func DisableAIDownloadActions() []DisableAIDownloadAction {
-	actions := make([]DisableAIDownloadAction, 0, len(AIDownloadFlagNames)+1)
-	for _, name := range AIDownloadFlagNames {
+// DisableAIDownloadActions returns one action per managed item — every entry
+// in AvailableAIDownloadFlags plus the Enterprise policy — describing
+// whether it will be applied (selected / includePolicy) or reverted to
+// Chrome's default (not selected / !includePolicy). The policy is
+// independent of which flags are selected.
+func DisableAIDownloadActions(selectedFlags []string, includePolicy bool) []DisableAIDownloadAction {
+	selected := make(map[string]bool, len(selectedFlags))
+	for _, name := range selectedFlags {
+		selected[name] = true
+	}
+
+	actions := make([]DisableAIDownloadAction, 0, len(AvailableAIDownloadFlags)+1)
+	for _, f := range AvailableAIDownloadFlags {
+		if selected[f.Name] {
+			actions = append(actions, DisableAIDownloadAction{
+				Label: "chrome://flags/#" + f.Name + " -> Disabled",
+			})
+			continue
+		}
 		actions = append(actions, DisableAIDownloadAction{
-			Label: "chrome://flags/#" + name + " -> Disabled",
+			Label:  "chrome://flags/#" + f.Name + " -> Default (reset)",
+			Revert: true,
+		})
+	}
+
+	if includePolicy {
+		actions = append(actions, DisableAIDownloadAction{
+			Label:            GenAIPolicyName + " = 1 (Disabled)",
+			Detail:           policyStorageDescription(true),
+			EnterprisePolicy: true,
+			PolicyNote:       `Chrome will show the "managed by your organization" banner`,
+		})
+	} else {
+		actions = append(actions, DisableAIDownloadAction{
+			Label:            GenAIPolicyName + " removed (reset to default)",
+			Detail:           policyStorageDescription(false),
+			EnterprisePolicy: true,
+			PolicyNote:       `Removes the "managed by your organization" banner, if shown`,
+			Revert:           true,
 		})
 	}
-	actions = append(actions, DisableAIDownloadAction{
-		Label:            GenAIPolicyName + " = 1 (Disabled)",
-		Detail:           policyStorageDescription(),
-		EnterprisePolicy: true,
-		PolicyNote:       `Chrome will show the "managed by your organization" banner`,
-	})
 	return actions
 }
 
+// GroupDisableAIDownloadActions splits actions (as returned by
+// DisableAIDownloadActions) into three buckets for display: chrome://flags
+// entries to apply, chrome://flags entries to revert, and the Enterprise
+// policy action (apply or revert).
+func GroupDisableAIDownloadActions(actions []DisableAIDownloadAction) (applyFlags, revertFlags, policy []DisableAIDownloadAction) {
+	for _, a := range actions {
+		switch {
+		case a.EnterprisePolicy:
+			policy = append(policy, a)
+		case a.Revert:
+			revertFlags = append(revertFlags, a)
+		default:
+			applyFlags = append(applyFlags, a)
+		}
+	}
+	return applyFlags, revertFlags, policy
+}
 
-// setFlagsDisabled rewrites browser.enabled_labs_experiments so each requested
-// flag appears exactly once with the Disabled choice (@2). Returns the list
-// of flag names whose state actually changed.
-func setFlagsDisabled(localState map[string]any, flags []string) []string {
+// syncManagedFlags rewrites browser.enabled_labs_experiments so every flag
+// in managed is Disabled (@2) when it also appears in selected, or has any
+// existing override removed (reverted to Chrome's default) when it does
+// not. Flags outside managed are left untouched. Returns the flags that
+// were newly disabled and the flags whose override was removed.
+func syncManagedFlags(localState map[string]any, managed, selected []string) (disabled, reverted []string) {
 	browser, _ := localState["browser"].(map[string]any)
 	if browser == nil {
 		browser = map[string]any{}
@@ -62,44 +123,57 @@ func setFlagsDisabled(localState map[string]any, flags []string) []string {
 		}
 	}
 
-	targets := make(map[string]bool, len(flags))
-	for _, name := range flags {
-		targets[name] = true
+	isManaged := make(map[string]bool, len(managed))
+	for _, name := range managed {
+		isManaged[name] = true
+	}
+	isSelected := make(map[string]bool, len(selected))
+	for _, name := range selected {
+		isSelected[name] = true
 	}
 
 	kept := make([]string, 0, len(existing))
-	alreadyDisabled := make(map[string]bool, len(flags))
+	alreadyDisabled := make(map[string]bool, len(selected))
+	wasPresent := make(map[string]bool, len(managed))
 	for _, entry := range existing {
 		name := entry
 		if idx := strings.IndexByte(entry, '@'); idx >= 0 {
 			name = entry[:idx]
 		}
-		if targets[name] {
-			if entry == name+flagDisabledSuffix && !alreadyDisabled[name] {
-				alreadyDisabled[name] = true
-				kept = append(kept, entry)
-			}
+		if !isManaged[name] {
+			kept = append(kept, entry)
 			continue
 		}
-		kept = append(kept, entry)
+		wasPresent[name] = true
+		if !isSelected[name] {
+			continue // revert: drop the existing override entirely
+		}
+		if entry == name+flagDisabledSuffix && !alreadyDisabled[name] {
+			alreadyDisabled[name] = true
+			kept = append(kept, entry)
+		}
 	}
 
-	changed := make([]string, 0, len(flags))
-	for _, name := range flags {
+	for _, name := range selected {
 		if !alreadyDisabled[name] {
 			kept = append(kept, name+flagDisabledSuffix)
-			changed = append(changed, name)
+			disabled = append(disabled, name)
+		}
+	}
+	for _, name := range managed {
+		if !isSelected[name] && wasPresent[name] {
+			reverted = append(reverted, name)
 		}
 	}
 
-	if len(changed) == 0 && len(kept) == len(existing) {
-		return nil
+	if len(disabled) == 0 && len(reverted) == 0 && len(kept) == len(existing) {
+		return nil, nil
 	}
 
 	next := make([]any, len(kept))
 	for i, s := range kept {
 		next[i] = s
 	}
 	browser["enabled_labs_experiments"] = next
-	return changed
+	return disabled, reverted
 }