diff --git a/cmd/aibomgen-cli/generate.go b/cmd/aibomgen-cli/generate.go index ed64e51..3bd569b 100644 --- a/cmd/aibomgen-cli/generate.go +++ b/cmd/aibomgen-cli/generate.go @@ -210,7 +210,7 @@ func runGenerate(cmd *cobra.Command, args []string) error { // Print summary. if len(written) == 0 { - genUI.PrintNoModelsFound() + genUI.PrintNoBOMsWritten() return nil } @@ -245,7 +245,7 @@ func runModelIDMode(genUI *ui.GenerateUI, modelIDs []string, mode, hfToken strin if !quiet { workflow = ui.NewWorkflow(os.Stdout, "") - processTaskIdx = workflow.AddTask("Processing models") + processTaskIdx = workflow.AddTask("Processing possible models") writeTaskIdx = workflow.AddTask("Writing output") workflow.Start() } @@ -327,7 +327,7 @@ func runModelIDMode(genUI *ui.GenerateUI, modelIDs []string, mode, hfToken strin } if !quiet && workflow != nil { - workflow.CompleteTask(processTaskIdx, fmt.Sprintf("%d model(s)", len(boms))) + workflow.CompleteTask(processTaskIdx, fmt.Sprintf("%d possible model(s)", len(modelIDs))) workflow.StartTask(writeTaskIdx, "") workflow.CompleteTask(writeTaskIdx, fmt.Sprintf("%d file(s)", len(boms))) workflow.Stop() @@ -389,14 +389,17 @@ type modelTracker struct { // detail is empty for a clean success (datasets are shown on sub-lines instead). func modelOutcome(t *modelTracker, hasToken bool) (mark, detail string) { switch { - case t == nil || !t.complete: + case t == nil: return ui.GetCrossMark(), ui.Error.Render("→ BOM build failed") - case !t.apiOK && t.notFound: + case t.notFound && !t.apiOK: if hasToken { - return ui.GetWarnMark(), ui.Warning.Render("→ not found on HF Hub") + return ui.GetCrossMark(), ui.Error.Render("→ not found on HF Hub; no BOM written") } - return ui.GetWarnMark(), ui.Warning.Render("→ not found on HF Hub (or private – set --hf-token)") + return ui.GetCrossMark(), ui.Error.Render("→ not found (or private – set --hf-token); no BOM written") + + case !t.complete: + return ui.GetCrossMark(), ui.Error.Render("→ BOM build failed") case t.fetchErr: if fetcher.IsUnauthorized(t.fetchErrVal) { diff --git a/cmd/aibomgen-cli/scan.go b/cmd/aibomgen-cli/scan.go index 0be4e44..148e7be 100644 --- a/cmd/aibomgen-cli/scan.go +++ b/cmd/aibomgen-cli/scan.go @@ -163,7 +163,7 @@ func runScan(cmd *cobra.Command, args []string) error { // Print summary. if len(written) == 0 { genUI := ui.NewGenerateUI(cmd.OutOrStdout(), quiet) - genUI.PrintNoModelsFound() + genUI.PrintNoBOMsWritten() return nil } @@ -202,8 +202,8 @@ func runScanDirectory(inputPath, mode, hfToken string, timeout time.Duration, qu if !quiet { workflow = ui.NewWorkflow(os.Stdout, "") - scanTaskIdx = workflow.AddTask("Scanning for AI imports") - processTaskIdx = workflow.AddTask("Processing models") + scanTaskIdx = workflow.AddTask("Scanning for possible AI imports") + processTaskIdx = workflow.AddTask("Processing possible models") writeTaskIdx = workflow.AddTask("Writing output") workflow.Start() } @@ -223,7 +223,7 @@ func runScanDirectory(inputPath, mode, hfToken string, timeout time.Duration, qu } if !quiet && workflow != nil { - workflow.CompleteTask(scanTaskIdx, fmt.Sprintf("found %d model(s)", len(discoveries))) + workflow.CompleteTask(scanTaskIdx, fmt.Sprintf("found %d possible model(s)", len(discoveries))) } if len(discoveries) == 0 { @@ -309,7 +309,7 @@ func runScanDirectory(inputPath, mode, hfToken string, timeout time.Duration, qu } if !quiet && workflow != nil { - workflow.CompleteTask(processTaskIdx, fmt.Sprintf("%d model(s)", len(boms))) + workflow.CompleteTask(processTaskIdx, fmt.Sprintf("%d possible model(s)", len(discoveries))) workflow.StartTask(writeTaskIdx, "") workflow.CompleteTask(writeTaskIdx, fmt.Sprintf("%d file(s)", len(boms))) workflow.Stop() diff --git a/internal/ui/generate.go b/internal/ui/generate.go index 1bf2970..61b449b 100644 --- a/internal/ui/generate.go +++ b/internal/ui/generate.go @@ -201,14 +201,14 @@ func (g *GenerateUI) PrintSummary(filesWritten int, outputDir, format string) { fmt.Fprintln(g.writer, SuccessBox.Render(summary.String())) } -// PrintNoModelsFound prints a message when no models are found. -func (g *GenerateUI) PrintNoModelsFound() { +// PrintNoBOMsWritten prints a message when no BOMs were written. +func (g *GenerateUI) PrintNoBOMsWritten() { if g.quiet { return } - msg := "No models detected; no AIBOM files written." - fmt.Fprintln(g.writer, Warning.Render(GetWarnMark()+" "+msg)) + msg := "No BOMs written." + fmt.Fprintln(g.writer, Error.Render(GetCrossMark()+" "+msg)) } // LogStep prints a simple log message (non-workflow mode). diff --git a/pkg/aibomgen/generator/generator.go b/pkg/aibomgen/generator/generator.go index 335f367..7387fb9 100644 --- a/pkg/aibomgen/generator/generator.go +++ b/pkg/aibomgen/generator/generator.go @@ -203,15 +203,25 @@ func BuildPerDiscovery(discoveries []scanner.Discovery, opts GenerateOptions) ([ var resp *fetcher.ModelAPIResponse var readme *fetcher.ModelReadmeCard + var apiNotFound bool if modelID != "" { if r, err := fetchers.modelAPI.Fetch(modelID); err == nil { resp = r progress(ProgressEvent{Type: EventFetchAPIComplete, ModelID: modelID}) } else { + if fetcher.IsNotFound(err) || fetcher.IsUnauthorized(err) { + apiNotFound = true + } progress(ProgressEvent{Type: EventError, ModelID: modelID, Error: err, Message: fetchErrMessage("API", err)}) } + // Skip BOM generation if API fetch returned not found or unauthorized (model not accessible on HF) + if apiNotFound { + progress(ProgressEvent{Type: EventModelComplete, ModelID: modelID, Message: "model skipped: API not found or unauthorized"}) + continue + } + if c, err := fetchers.modelReadme.Fetch(modelID); err == nil { readme = c progress(ProgressEvent{Type: EventFetchReadmeComplete, ModelID: modelID}) @@ -385,19 +395,29 @@ func BuildFromModelIDs(modelIDs []string, opts GenerateOptions) ([]DiscoveredBOM continue } - bomBuilder := newBOMBuilder() - progress(ProgressEvent{Type: EventFetchStart, ModelID: modelID, Index: i, Total: len(modelIDs)}) // Fetch API metadata. resp, err := fetchers.modelAPI.Fetch(modelID) + var apiNotFound bool if err != nil { + if fetcher.IsNotFound(err) || fetcher.IsUnauthorized(err) { + apiNotFound = true + } progress(ProgressEvent{Type: EventError, ModelID: modelID, Error: err, Message: "API fetch failed"}) resp = nil } else { progress(ProgressEvent{Type: EventFetchAPIComplete, ModelID: modelID}) } + // Skip BOM generation if API fetch returned not found or unauthorized (model not accessible on HF) + if apiNotFound { + progress(ProgressEvent{Type: EventModelComplete, ModelID: modelID, Message: "model skipped: API not found or unauthorized"}) + continue + } + + bomBuilder := newBOMBuilder() + // Fetch README. readme, err := fetchers.modelReadme.Fetch(modelID) if err != nil { diff --git a/pkg/aibomgen/generator/generator_test.go b/pkg/aibomgen/generator/generator_test.go index c6d603b..6568330 100644 --- a/pkg/aibomgen/generator/generator_test.go +++ b/pkg/aibomgen/generator/generator_test.go @@ -78,6 +78,32 @@ func (m *mockDatasetReadmeFetcher) Fetch(id string) (*fetcher.DatasetReadmeCard, return &fetcher.DatasetReadmeCard{}, nil } +func successFetcherSet() fetcherSet { + return fetcherSet{ + modelAPI: &mockModelAPIFetcher{ + fetchFunc: func(id string) (*fetcher.ModelAPIResponse, error) { + return &fetcher.ModelAPIResponse{ID: id}, nil + }, + }, + modelReadme: &mockModelReadmeFetcher{ + fetchFunc: func(id string) (*fetcher.ModelReadmeCard, error) { + return &fetcher.ModelReadmeCard{}, nil + }, + }, + datasetAPI: &mockDatasetAPIFetcher{ + fetchFunc: func(id string) (*fetcher.DatasetAPIResponse, error) { + return &fetcher.DatasetAPIResponse{ID: id}, nil + }, + }, + datasetReadme: &mockDatasetReadmeFetcher{ + fetchFunc: func(id string) (*fetcher.DatasetReadmeCard, error) { + return &fetcher.DatasetReadmeCard{}, nil + }, + }, + modelTree: &fetcher.DummyModelTreeFetcher{}, + } +} + func TestBuildDummyBOM(t *testing.T) { // Save originals. originalDummyFetcherSet := newDummyFetcherSet @@ -271,6 +297,9 @@ func TestBuildPerDiscovery(t *testing.T) { }, } } + newFetcherSet = func(httpClient *http.Client) fetcherSet { + return successFetcherSet() + } }, wantErr: false, check: func(t *testing.T, got []DiscoveredBOM) { @@ -311,6 +340,9 @@ func TestBuildPerDiscovery(t *testing.T) { }, } } + newFetcherSet = func(httpClient *http.Client) fetcherSet { + return successFetcherSet() + } }, wantErr: false, check: func(t *testing.T, got []DiscoveredBOM) { @@ -335,6 +367,9 @@ func TestBuildPerDiscovery(t *testing.T) { }, } } + newFetcherSet = func(httpClient *http.Client) fetcherSet { + return successFetcherSet() + } }, wantErr: false, check: func(t *testing.T, got []DiscoveredBOM) { @@ -362,6 +397,9 @@ func TestBuildPerDiscovery(t *testing.T) { }, } } + newFetcherSet = func(httpClient *http.Client) fetcherSet { + return successFetcherSet() + } }, wantErr: false, check: func(t *testing.T, got []DiscoveredBOM) { @@ -613,6 +651,42 @@ func TestBuildPerDiscovery(t *testing.T) { } }, }, + { + name: "skips model when API returns 401 (unauthorized) (model not found is also 401)", + args: args{ + discoveries: []scanner.Discovery{ + {ID: "private-or-missing-model", Name: "private-or-missing-model", Type: "huggingface"}, + }, + opts: GenerateOptions{Timeout: 1 * time.Second}, + }, + setup: func() { + newBOMBuilder = func() bomBuilder { + return &mockBOMBuilder{ + buildFunc: func(bctx builder.BuildContext) (*cdx.BOM, error) { + return &cdx.BOM{}, nil + }, + } + } + newFetcherSet = func(httpClient *http.Client) fetcherSet { + return fetcherSet{ + modelAPI: &mockModelAPIFetcher{ + fetchFunc: func(id string) (*fetcher.ModelAPIResponse, error) { + return nil, &fetcher.HFError{StatusCode: 401} + }, + }, + modelReadme: &mockModelReadmeFetcher{}, + datasetAPI: &mockDatasetAPIFetcher{}, + datasetReadme: &mockDatasetReadmeFetcher{}, + } + } + }, + wantErr: false, + check: func(t *testing.T, got []DiscoveredBOM) { + if len(got) != 0 { + t.Errorf("Expected 0 BOMs for model unauthorized on HF (401), got %d", len(got)) + } + }, + }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { @@ -766,6 +840,9 @@ func TestBuildFromModelIDs(t *testing.T) { }, } } + newFetcherSet = func(httpClient *http.Client) fetcherSet { + return successFetcherSet() + } }, wantErr: false, check: func(t *testing.T, got []DiscoveredBOM) { @@ -792,6 +869,9 @@ func TestBuildFromModelIDs(t *testing.T) { }, } } + newFetcherSet = func(httpClient *http.Client) fetcherSet { + return successFetcherSet() + } }, wantErr: false, check: func(t *testing.T, got []DiscoveredBOM) { @@ -833,6 +913,9 @@ func TestBuildFromModelIDs(t *testing.T) { }, } } + newFetcherSet = func(httpClient *http.Client) fetcherSet { + return successFetcherSet() + } }, wantErr: false, check: func(t *testing.T, got []DiscoveredBOM) { @@ -860,6 +943,9 @@ func TestBuildFromModelIDs(t *testing.T) { }, } } + newFetcherSet = func(httpClient *http.Client) fetcherSet { + return successFetcherSet() + } }, wantErr: false, check: func(t *testing.T, got []DiscoveredBOM) { @@ -982,6 +1068,74 @@ func TestBuildFromModelIDs(t *testing.T) { } }, }, + { + name: "skips model when API returns 404 (not found)", + args: args{ + modelIDs: []string{"org/nonexistent-model"}, + opts: GenerateOptions{Timeout: 1 * time.Second}, + }, + setup: func() { + newBOMBuilder = func() bomBuilder { + return &mockBOMBuilder{ + buildFunc: func(bctx builder.BuildContext) (*cdx.BOM, error) { + return &cdx.BOM{}, nil + }, + } + } + newFetcherSet = func(httpClient *http.Client) fetcherSet { + return fetcherSet{ + modelAPI: &mockModelAPIFetcher{ + fetchFunc: func(id string) (*fetcher.ModelAPIResponse, error) { + return nil, &fetcher.HFError{StatusCode: 404} + }, + }, + modelReadme: &mockModelReadmeFetcher{}, + datasetAPI: &mockDatasetAPIFetcher{}, + datasetReadme: &mockDatasetReadmeFetcher{}, + } + } + }, + wantErr: false, + check: func(t *testing.T, got []DiscoveredBOM) { + if len(got) != 0 { + t.Errorf("Expected 0 BOMs for model not found on HF (404), got %d", len(got)) + } + }, + }, + { + name: "skips model when API returns 401 (unauthorized)", + args: args{ + modelIDs: []string{"org/private-or-missing-model"}, + opts: GenerateOptions{Timeout: 1 * time.Second}, + }, + setup: func() { + newBOMBuilder = func() bomBuilder { + return &mockBOMBuilder{ + buildFunc: func(bctx builder.BuildContext) (*cdx.BOM, error) { + return &cdx.BOM{}, nil + }, + } + } + newFetcherSet = func(httpClient *http.Client) fetcherSet { + return fetcherSet{ + modelAPI: &mockModelAPIFetcher{ + fetchFunc: func(id string) (*fetcher.ModelAPIResponse, error) { + return nil, &fetcher.HFError{StatusCode: 401} + }, + }, + modelReadme: &mockModelReadmeFetcher{}, + datasetAPI: &mockDatasetAPIFetcher{}, + datasetReadme: &mockDatasetReadmeFetcher{}, + } + } + }, + wantErr: false, + check: func(t *testing.T, got []DiscoveredBOM) { + if len(got) != 0 { + t.Errorf("Expected 0 BOMs for model unauthorized on HF (401), got %d", len(got)) + } + }, + }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) {