diff --git a/.github/workflows/agent-build.yml b/.github/workflows/agent-build.yml new file mode 100644 index 000000000..9da25b73a --- /dev/null +++ b/.github/workflows/agent-build.yml @@ -0,0 +1,84 @@ +# Copyright 2025 RAIDS Lab +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +name: Agent Build + +on: + push: + branches: [main] + paths: + - "crater-agent/**" + - ".github/workflows/agent-build.yml" + tags: + - "v*.*.*" + +env: + REGISTRY: ghcr.io + REPOSITORY: raids-lab + IMAGE_NAME: crater-agent + +jobs: + build-and-push-image: + runs-on: ubuntu-latest + permissions: + contents: read + packages: write + steps: + - name: Checkout code + uses: actions/checkout@v5 + + - name: Set up QEMU + uses: docker/setup-qemu-action@v3 + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + + - name: Login to GHCR + uses: docker/login-action@v3 + with: + registry: ${{ env.REGISTRY }} + username: ${{ github.repository_owner }} + password: ${{ secrets.GITHUB_TOKEN }} + + - name: Docker meta + id: meta + uses: docker/metadata-action@v5 + with: + images: ${{ env.REGISTRY }}/${{ env.REPOSITORY }}/${{ env.IMAGE_NAME }} + tags: | + type=ref,event=branch + type=semver,pattern={{version}} + type=semver,pattern={{major}}.{{minor}} + type=semver,pattern={{major}} + type=raw,value=latest,enable={{is_default_branch}} + type=sha + + - name: Build and push multi-platform image + uses: docker/build-push-action@v6 + with: + context: ./crater-agent + file: ./crater-agent/Dockerfile + platforms: linux/amd64,linux/arm64 + push: true + tags: ${{ steps.meta.outputs.tags }} + + - uses: quartx-analytics/ghcr-cleaner@v1 + with: + owner-type: org + token: ${{ secrets.PAT_TOKEN }} + repository-owner: ${{ github.repository_owner }} + package-name: ${{ env.IMAGE_NAME }} + delete-untagged: true + keep-at-most: 2 + skip-tags: v* diff --git a/.github/workflows/agent-pr.yml b/.github/workflows/agent-pr.yml new file mode 100644 index 000000000..07c10f498 --- /dev/null +++ b/.github/workflows/agent-pr.yml @@ -0,0 +1,42 @@ +# Copyright 2025 RAIDS Lab +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +name: Agent PR Check + +on: + pull_request: + branches: [main] + paths: + - "crater-agent/**" + - ".github/workflows/agent-pr.yml" + +jobs: + agent-build-check: + runs-on: ubuntu-latest + permissions: + contents: read + steps: + - name: Checkout code + uses: actions/checkout@v5 + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + + - name: Build agent image + uses: docker/build-push-action@v6 + with: + context: ./crater-agent + file: ./crater-agent/Dockerfile + push: false + tags: crater-agent:pr-check diff --git a/.github/workflows/backend-build.yml b/.github/workflows/backend-build.yml index 7c5f5a7d9..5e81068b4 100644 --- a/.github/workflows/backend-build.yml +++ b/.github/workflows/backend-build.yml @@ -130,7 +130,6 @@ jobs: mkdir -p bin/${{ matrix.platform.image_platform }} go build -ldflags="-X main.AppVersion=${{ steps.set-version.outputs.app_version }} -X main.CommitSHA=${{ steps.set-version.outputs.commit_sha }} -X main.BuildType=${{ steps.set-version.outputs.build_type }} -X main.BuildTime=${{ steps.set-version.outputs.build_time }} -w -s" -o bin/${{ matrix.platform.image_platform }}/migrate cmd/gorm-gen/models/migrate.go go build -ldflags="-X main.AppVersion=${{ steps.set-version.outputs.app_version }} -X main.CommitSHA=${{ steps.set-version.outputs.commit_sha }} -X main.BuildType=${{ steps.set-version.outputs.build_type }} -X main.BuildTime=${{ steps.set-version.outputs.build_time }} -w -s" -o bin/${{ matrix.platform.image_platform }}/controller cmd/crater/main.go - go build -ldflags="-w -s" -o bin/${{ matrix.platform.image_platform }}/model-metadata-refresh cmd/model-metadata-refresh/main.go env: CGO_ENABLED: 0 GOOS: ${{ matrix.platform.goos }} diff --git a/.github/workflows/storage-build.yml b/.github/workflows/storage-build.yml index ae3ed338a..be3be5f65 100644 --- a/.github/workflows/storage-build.yml +++ b/.github/workflows/storage-build.yml @@ -20,11 +20,6 @@ on: paths: - "backend/internal/storage/**" - "backend/cmd/storage-server/**" - - "backend/cmd/model-dataset-governance/**" - - "backend/internal/governance/modeldataset/**" - - "backend/dao/model/**" - - "backend/dao/query/**" - - "backend/pkg/config/**" - "backend/storage-server.Dockerfile" - "backend/go.mod" - "backend/go.sum" @@ -134,7 +129,6 @@ jobs: run: | mkdir -p bin/${{ matrix.platform.image_platform }} go build -ldflags="-X main.AppVersion=${{ steps.set-version.outputs.app_version }} -X main.CommitSHA=${{ steps.set-version.outputs.commit_sha }} -X main.BuildType=${{ steps.set-version.outputs.build_type }} -X main.BuildTime=${{ steps.set-version.outputs.build_time }} -w -s" -o bin/${{ matrix.platform.image_platform }}/storage-server cmd/storage-server/main.go - go build -ldflags="-X main.AppVersion=${{ steps.set-version.outputs.app_version }} -X main.CommitSHA=${{ steps.set-version.outputs.commit_sha }} -X main.BuildType=${{ steps.set-version.outputs.build_type }} -X main.BuildTime=${{ steps.set-version.outputs.build_time }} -w -s" -o bin/${{ matrix.platform.image_platform }}/model-dataset-governance cmd/model-dataset-governance/main.go env: CGO_ENABLED: 0 GOOS: ${{ matrix.platform.goos }} diff --git a/.gitignore b/.gitignore index 83c85888e..948da401f 100644 --- a/.gitignore +++ b/.gitignore @@ -40,6 +40,7 @@ dist-ssr *.dll *.so *.dylib +backend/crater bin testbin/* Dockerfile.cross @@ -51,6 +52,7 @@ etc/* # Configuration backups (created by make config-link) *.bak +crater-agent/config/*backup*.json # Test binary, build with `go test -c` *.test @@ -79,3 +81,5 @@ kubeconfig* # Claude Code .claude/ +storage/etc/config.yaml +.codex diff --git a/Makefile b/Makefile index 06d5bd333..188f39f76 100644 --- a/Makefile +++ b/Makefile @@ -57,6 +57,33 @@ config-unlink: ## Remove configuration symlinks (only symlinks, not regular file config-restore: ## Restore configuration files from .bak backups @bash hack/config.sh restore +##@ Development + +.PHONY: run-agent +run-agent: ## Run the local Python crater-agent service. + @echo "$(GREEN)Starting crater-agent on http://localhost:8000 ...$(RESET)" + @cd crater-agent && \ + if [ ! -d .venv ]; then python3.11 -m venv .venv || python3 -m venv .venv; fi && \ + .venv/bin/python -m pip install -e . && \ + .venv/bin/python -m uvicorn crater_agent.app:app --host 0.0.0.0 --port 8000 + +.PHONY: run-backend +run-backend: ## Run the local Go backend. + @$(MAKE) -C backend run + +.PHONY: run-frontend +run-frontend: ## Run the local React frontend. + @$(MAKE) -C frontend run + +.PHONY: run-dev +run-dev: ## Run agent, backend, and frontend together for local development. + @echo "$(GREEN)Starting crater agent, backend, and frontend. Press Ctrl-C to stop all.$(RESET)" + @set -e; \ + trap 'for pid in $$(jobs -p); do kill $$pid 2>/dev/null || true; done' INT TERM EXIT; \ + $(MAKE) run-agent & \ + $(MAKE) run-backend & \ + $(MAKE) run-frontend & \ + wait + # 默认目标 .DEFAULT_GOAL := help - diff --git a/backend/.gitlab-ci.yml b/backend/.gitlab-ci.yml index c6ee93eaa..058f54905 100644 --- a/backend/.gitlab-ci.yml +++ b/backend/.gitlab-ci.yml @@ -48,7 +48,6 @@ build_binaries: - go run github.com/swaggo/swag/cmd/swag@latest init - go build -ldflags="-w -s" -o "${BIN_DIR}/migrate" cmd/gorm-gen/models/migrate.go - go build -ldflags="-w -s" -o "${BIN_DIR}/controller" main.go - - go build -ldflags="-w -s" -o "${BIN_DIR}/model-metadata-refresh" cmd/model-metadata-refresh/main.go artifacts: paths: - "${BIN_DIR}" diff --git a/backend/Dockerfile b/backend/Dockerfile index cd94d4d5d..1a86c581d 100644 --- a/backend/Dockerfile +++ b/backend/Dockerfile @@ -30,8 +30,7 @@ ENV GIN_MODE=release COPY LICENSE NOTICE /usr/share/doc/crater/ COPY $BIN_DIR/bin-${TARGETPLATFORM//\//_}/controller . COPY $BIN_DIR/bin-${TARGETPLATFORM//\//_}/migrate . -COPY $BIN_DIR/bin-${TARGETPLATFORM//\//_}/model-metadata-refresh . -RUN chmod +x controller migrate model-metadata-refresh +RUN chmod +x controller migrate EXPOSE 8088:8088 diff --git a/backend/Makefile b/backend/Makefile index 60ef69f61..0e4aee284 100644 --- a/backend/Makefile +++ b/backend/Makefile @@ -192,12 +192,6 @@ build-migrate: fmt lint ## Build migration binary. CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build -o bin/migrate cmd/gorm-gen/models/migrate.go @echo "$(GREEN)✅ Migration binary built successfully: bin/migrate$(RESET)" -.PHONY: build-model-metadata-refresh -build-model-metadata-refresh: fmt lint ## Build model metadata refresh binary. - @echo "$(YELLOW)Building model metadata refresh binary...$(RESET)" - CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build -o bin/model-metadata-refresh cmd/model-metadata-refresh/main.go - @echo "$(GREEN)Model metadata refresh binary built successfully: bin/model-metadata-refresh$(RESET)" - ##@ Development Tools ## Location to install tools to @@ -240,5 +234,4 @@ run-storage: prepare fmt imports ## Run storage server from backend module. build-storage: fmt lint ## Build storage server binary. @echo "$(YELLOW)Building storage server binary...$(RESET)" go build -ldflags="-w -s $(VERSION_LDFLAGS)" -o bin/storage-server cmd/storage-server/main.go - go build -ldflags="-w -s $(VERSION_LDFLAGS)" -o bin/model-dataset-governance cmd/model-dataset-governance/main.go @echo "$(GREEN)Storage server binary built successfully: bin/storage-server$(RESET)" diff --git a/backend/cmd/gorm-gen/curd/generate.go b/backend/cmd/gorm-gen/curd/generate.go index 3fa17cb46..9639e5d23 100644 --- a/backend/cmd/gorm-gen/curd/generate.go +++ b/backend/cmd/gorm-gen/curd/generate.go @@ -26,7 +26,6 @@ func main() { model.User{}, model.Account{}, model.UserAccount{}, - model.ModelDatasetSource{}, model.Dataset{}, model.AccountDataset{}, model.UserDataset{}, @@ -46,8 +45,6 @@ func main() { model.CronJobRecord{}, model.CronJobConfig{}, model.ModelDownload{}, - model.ModelDownloadSubmission{}, - model.ModelDatasetDiscovery{}, model.UserModelDownload{}, model.GpuAnalysis{}, model.SystemConfig{}, diff --git a/backend/cmd/gorm-gen/models/migrate.go b/backend/cmd/gorm-gen/models/migrate.go index 53a9c32f5..c96749f30 100644 --- a/backend/cmd/gorm-gen/models/migrate.go +++ b/backend/cmd/gorm-gen/models/migrate.go @@ -33,177 +33,6 @@ type modelDownloadSourceMetadataMigration struct { SourceCreatedAt *time.Time `gorm:"comment:源站创建时间"` } -func modelDatasetSourceMigration() *gormigrate.Migration { - return &gormigrate.Migration{ - ID: "202607121200", - Migrate: func(tx *gorm.DB) error { - if err := createTableIfMissing(tx, &model.ModelDatasetSource{}); err != nil { - return err - } - if err := createTableIfMissing(tx, &model.ModelDatasetDiscovery{}); err != nil { - return err - } - if err := addColumnIfMissing(tx, "datasets", &model.Dataset{}, "ModelDatasetSourceID"); err != nil { - return err - } - if err := addColumnIfMissing( - tx, "model_downloads", &model.ModelDownload{}, "ModelDatasetSourceID", - ); err != nil { - return err - } - if err := createConstraintIfMissing( - tx, "datasets", &model.Dataset{}, "ModelDatasetSource", - ); err != nil { - return err - } - if err := createConstraintIfMissing( - tx, "model_downloads", &model.ModelDownload{}, "ModelDatasetSource", - ); err != nil { - return err - } - if err := createIndexIfMissing( - tx, "datasets", &model.Dataset{}, "ModelDatasetSourceID", - ); err != nil { - return err - } - return createIndexIfMissing( - tx, "model_downloads", &model.ModelDownload{}, "ModelDatasetSourceID", - ) - }, - Rollback: func(tx *gorm.DB) error { - if err := dropIndexIfPresent( - tx, "model_downloads", &model.ModelDownload{}, "ModelDatasetSourceID", - ); err != nil { - return err - } - if err := dropIndexIfPresent(tx, "datasets", &model.Dataset{}, "ModelDatasetSourceID"); err != nil { - return err - } - if err := dropConstraintIfPresent( - tx, "model_downloads", &model.ModelDownload{}, "ModelDatasetSource", - ); err != nil { - return err - } - if err := dropConstraintIfPresent( - tx, "datasets", &model.Dataset{}, "ModelDatasetSource", - ); err != nil { - return err - } - if err := dropColumnIfPresent( - tx, "model_downloads", &model.ModelDownload{}, "ModelDatasetSourceID", - ); err != nil { - return err - } - if err := dropColumnIfPresent(tx, "datasets", &model.Dataset{}, "ModelDatasetSourceID"); err != nil { - return err - } - if err := dropTableIfPresent(tx, &model.ModelDatasetDiscovery{}); err != nil { - return err - } - return dropTableIfPresent(tx, &model.ModelDatasetSource{}) - }, - } -} - -func modelDownloadSubmissionMigration() *gormigrate.Migration { - return &gormigrate.Migration{ - ID: "202607181200", - Migrate: func(tx *gorm.DB) error { - if err := createTableIfMissing(tx, &model.ModelDownloadSubmission{}); err != nil { - return err - } - if !tx.Migrator().HasTable(&model.ModelDownload{}) { - return nil - } - // Active downloads that predate quota tracking still need one owner - // reservation so an upgrade cannot temporarily bypass both limits. - return tx.Exec(` - INSERT INTO model_download_submissions - (user_id, model_download_id, action, status, created_at) - SELECT download.creator_id, download.id, ?, ?, CURRENT_TIMESTAMP - FROM model_downloads AS download - WHERE download.status IN ? AND download.deleted_at IS NULL - AND NOT EXISTS ( - SELECT 1 FROM model_download_submissions AS submission - WHERE submission.model_download_id = download.id - AND submission.status = ? - )`, - model.ModelDownloadSubmissionCreate, - model.ModelDownloadSubmissionReserved, - []model.ModelDownloadStatus{ - model.ModelDownloadStatusPending, model.ModelDownloadStatusDownloading, - }, - model.ModelDownloadSubmissionReserved, - ).Error - }, - Rollback: func(tx *gorm.DB) error { - return dropTableIfPresent(tx, &model.ModelDownloadSubmission{}) - }, - } -} - -func createTableIfMissing(db *gorm.DB, value any) error { - if db.Migrator().HasTable(value) { - return nil - } - return db.Migrator().CreateTable(value) -} - -func dropTableIfPresent(db *gorm.DB, value any) error { - if !db.Migrator().HasTable(value) { - return nil - } - return db.Migrator().DropTable(value) -} - -func addColumnIfMissing(db *gorm.DB, table string, value any, field string) error { - migrator := db.Table(table).Migrator() - if migrator.HasColumn(value, field) { - return nil - } - return migrator.AddColumn(value, field) -} - -func dropColumnIfPresent(db *gorm.DB, table string, value any, field string) error { - migrator := db.Table(table).Migrator() - if !migrator.HasColumn(value, field) { - return nil - } - return migrator.DropColumn(value, field) -} - -func createConstraintIfMissing(db *gorm.DB, table string, value any, name string) error { - migrator := db.Table(table).Migrator() - if migrator.HasConstraint(value, name) { - return nil - } - return migrator.CreateConstraint(value, name) -} - -func dropConstraintIfPresent(db *gorm.DB, table string, value any, name string) error { - migrator := db.Table(table).Migrator() - if !migrator.HasConstraint(value, name) { - return nil - } - return migrator.DropConstraint(value, name) -} - -func createIndexIfMissing(db *gorm.DB, table string, value any, name string) error { - migrator := db.Table(table).Migrator() - if migrator.HasIndex(value, name) { - return nil - } - return migrator.CreateIndex(value, name) -} - -func dropIndexIfPresent(db *gorm.DB, table string, value any, name string) error { - migrator := db.Table(table).Migrator() - if !migrator.HasIndex(value, name) { - return nil - } - return migrator.DropIndex(value, name) -} - //nolint:gocyclo // ignore cyclomatic complexity func main() { db := query.GetDB() @@ -941,6 +770,119 @@ func main() { return tx.Migrator().DropTable(&model.GpuAnalysis{}, &model.SystemConfig{}) }, }, + { + ID: "202603290001", + Migrate: func(tx *gorm.DB) error { + return tx.AutoMigrate( + &model.AgentSession{}, + &model.AgentMessage{}, + &model.AgentToolCall{}, + &model.JobLogSnapshot{}, + ) + }, + Rollback: func(tx *gorm.DB) error { + return tx.Migrator().DropTable( + "agent_sessions", + "agent_messages", + "agent_tool_calls", + "job_log_snapshots", + ) + }, + }, + { + ID: "202604030001", + Migrate: func(tx *gorm.DB) error { + type AgentSession struct { + PinnedAt *time.Time `gorm:"index"` + } + return tx.Migrator().AddColumn(&AgentSession{}, "PinnedAt") + }, + Rollback: func(tx *gorm.DB) error { + type AgentSession struct { + PinnedAt *time.Time `gorm:"index"` + } + return tx.Migrator().DropColumn(&AgentSession{}, "PinnedAt") + }, + }, + { + ID: "202604030002", + Migrate: func(tx *gorm.DB) error { + if err := tx.AutoMigrate( + &model.AgentSession{}, + &model.AgentToolCall{}, + &model.AgentTurn{}, + &model.AgentRunEvent{}, + ); err != nil { + return err + } + return tx.Exec(` + UPDATE agent_sessions + SET last_orchestration_mode = 'single_agent' + WHERE last_orchestration_mode IS NULL + OR BTRIM(last_orchestration_mode) = '' + `).Error + }, + Rollback: func(tx *gorm.DB) error { + type AgentSession struct { + LastOrchestrationMode string `gorm:"type:varchar(32);default:'single_agent'"` + } + type AgentToolCall struct { + TurnID string `gorm:"type:uuid;index"` + ToolCallID string `gorm:"type:varchar(128);index"` + AgentID string `gorm:"type:varchar(128);index"` + ParentEventID *uint `gorm:"index"` + AgentRole string `gorm:"type:varchar(32);index"` + } + if err := tx.Migrator().DropTable("agent_run_events", "agent_turns"); err != nil { + return err + } + if err := tx.Migrator().DropColumn(&AgentToolCall{}, "AgentRole"); err != nil { + return err + } + if err := tx.Migrator().DropColumn(&AgentToolCall{}, "ParentEventID"); err != nil { + return err + } + if err := tx.Migrator().DropColumn(&AgentToolCall{}, "AgentID"); err != nil { + return err + } + if err := tx.Migrator().DropColumn(&AgentToolCall{}, "ToolCallID"); err != nil { + return err + } + if err := tx.Migrator().DropColumn(&AgentToolCall{}, "TurnID"); err != nil { + return err + } + return tx.Migrator().DropColumn(&AgentSession{}, "LastOrchestrationMode") + }, + }, + { + ID: "202604110001", + Migrate: func(tx *gorm.DB) error { + return runStatements(tx, []string{ + `ALTER TABLE agent_tool_calls + ADD COLUMN IF NOT EXISTS execution_backend VARCHAR(64), + ADD COLUMN IF NOT EXISTS sandbox_job_name VARCHAR(255), + ADD COLUMN IF NOT EXISTS script_name VARCHAR(128), + ADD COLUMN IF NOT EXISTS result_artifact_ref TEXT, + ADD COLUMN IF NOT EXISTS egress_domains JSONB`, + `CREATE INDEX IF NOT EXISTS idx_agent_tool_calls_sandbox_job_name + ON agent_tool_calls (sandbox_job_name)`, + `CREATE INDEX IF NOT EXISTS idx_agent_tool_calls_script_name + ON agent_tool_calls (script_name)`, + }) + }, + Rollback: func(tx *gorm.DB) error { + return runStatements(tx, []string{ + `DROP INDEX IF EXISTS idx_agent_tool_calls_script_name`, + `DROP INDEX IF EXISTS idx_agent_tool_calls_sandbox_job_name`, + `ALTER TABLE agent_tool_calls + DROP COLUMN IF EXISTS egress_domains, + DROP COLUMN IF EXISTS result_artifact_ref, + DROP COLUMN IF EXISTS script_name, + DROP COLUMN IF EXISTS sandbox_job_name, + DROP COLUMN IF EXISTS execution_backend`, + }) + }, + }, { ID: "202512261300", Migrate: func(tx *gorm.DB) error { @@ -1293,6 +1235,65 @@ func main() { return tx.Migrator().DropColumn(&Dataset{}, "MountCount") }, }, + { + ID: "202604220001", + Migrate: func(tx *gorm.DB) error { + return tx.AutoMigrate(&model.AgentFeedback{}) + }, + Rollback: func(tx *gorm.DB) error { + return tx.Migrator().DropTable("agent_feedbacks") + }, + }, + { + ID: "202604220002", + Migrate: func(tx *gorm.DB) error { + return runStatements(tx, []string{ + `ALTER TABLE agent_sessions + ADD COLUMN IF NOT EXISTS source VARCHAR(32) NOT NULL DEFAULT 'chat'`, + `UPDATE agent_sessions + SET source = 'chat' + WHERE source IS NULL OR BTRIM(source) = ''`, + `UPDATE agent_sessions + SET source = 'ops_audit' + WHERE title LIKE '[audit] 审批%' + AND source = 'chat'`, + `CREATE INDEX IF NOT EXISTS idx_agent_sessions_source + ON agent_sessions (source)`, + `ALTER TABLE agent_tool_calls + ADD COLUMN IF NOT EXISTS source VARCHAR(32) NOT NULL DEFAULT 'backend'`, + `UPDATE agent_tool_calls + SET source = 'backend' + WHERE source IS NULL OR BTRIM(source) = ''`, + `CREATE INDEX IF NOT EXISTS idx_agent_tool_calls_source + ON agent_tool_calls (source)`, + }) + }, + Rollback: func(tx *gorm.DB) error { + return runStatements(tx, []string{ + `DROP INDEX IF EXISTS idx_agent_tool_calls_source`, + `ALTER TABLE agent_tool_calls + DROP COLUMN IF EXISTS source`, + `DROP INDEX IF EXISTS idx_agent_sessions_source`, + `ALTER TABLE agent_sessions + DROP COLUMN IF EXISTS source`, + }) + }, + }, + { + ID: "202604230001", + Migrate: func(tx *gorm.DB) error { + type AgentFeedback struct { + EnrichedAt *time.Time `json:"enrichedAt,omitempty"` + } + return tx.Migrator().AddColumn(&AgentFeedback{}, "EnrichedAt") + }, + Rollback: func(tx *gorm.DB) error { + type AgentFeedback struct { + EnrichedAt *time.Time `json:"enrichedAt,omitempty"` + } + return tx.Migrator().DropColumn(&AgentFeedback{}, "EnrichedAt") + }, + }, { ID: "202604261000", Migrate: func(tx *gorm.DB) error { @@ -1324,6 +1325,33 @@ func main() { return nil }, }, + { + ID: "202604270001", + Migrate: func(tx *gorm.DB) error { + return runStatements(tx, []string{ + `DROP INDEX IF EXISTS idx_agent_tool_calls_script_name`, + `DROP INDEX IF EXISTS idx_agent_tool_calls_sandbox_job_name`, + `ALTER TABLE agent_tool_calls + DROP COLUMN IF EXISTS egress_domains, + DROP COLUMN IF EXISTS result_artifact_ref, + DROP COLUMN IF EXISTS script_name, + DROP COLUMN IF EXISTS sandbox_job_name`, + }) + }, + Rollback: func(tx *gorm.DB) error { + return runStatements(tx, []string{ + `ALTER TABLE agent_tool_calls + ADD COLUMN IF NOT EXISTS sandbox_job_name VARCHAR(255), + ADD COLUMN IF NOT EXISTS script_name VARCHAR(128), + ADD COLUMN IF NOT EXISTS result_artifact_ref TEXT, + ADD COLUMN IF NOT EXISTS egress_domains JSONB`, + `CREATE INDEX IF NOT EXISTS idx_agent_tool_calls_sandbox_job_name + ON agent_tool_calls (sandbox_job_name)`, + `CREATE INDEX IF NOT EXISTS idx_agent_tool_calls_script_name + ON agent_tool_calls (script_name)`, + }) + }, + }, { ID: "202607100100", Migrate: func(tx *gorm.DB) error { @@ -1381,7 +1409,9 @@ func main() { SourceUpdatedAt *time.Time `gorm:"comment:源站更新时间"` } migrator := tx.Table("model_downloads").Migrator() - for _, field := range []string{"Organization", "LogoURL", "SourceURL", "SourceDownloads", "SourceLikes", "SourceUpdatedAt"} { + for _, field := range []string{ + "Organization", "LogoURL", "SourceURL", "SourceDownloads", "SourceLikes", "SourceUpdatedAt", + } { if err := migrator.AddColumn(&ModelDownload{}, field); err != nil { return err } @@ -1398,7 +1428,9 @@ func main() { SourceUpdatedAt *time.Time } migrator := tx.Table("model_downloads").Migrator() - for _, field := range []string{"SourceUpdatedAt", "SourceLikes", "SourceDownloads", "SourceURL", "LogoURL", "Organization"} { + for _, field := range []string{ + "SourceUpdatedAt", "SourceLikes", "SourceDownloads", "SourceURL", "LogoURL", "Organization", + } { if err := migrator.DropColumn(&ModelDownload{}, field); err != nil { return err } @@ -1498,8 +1530,6 @@ func main() { ) }, }, - modelDatasetSourceMigration(), - modelDownloadSubmissionMigration(), }) m.InitSchema(func(tx *gorm.DB) error { @@ -1507,7 +1537,6 @@ func main() { &model.User{}, &model.Account{}, &model.UserAccount{}, - &model.ModelDatasetSource{}, &model.Dataset{}, &model.AccountDataset{}, &model.UserDataset{}, @@ -1525,9 +1554,7 @@ func main() { &model.ResourceNetwork{}, &model.ResourceVGPU{}, &model.ModelDownload{}, - &model.ModelDatasetDiscovery{}, &model.UserModelDownload{}, - &model.ModelDownloadSubmission{}, &model.CronJobConfig{}, &model.CronJobRecord{}, &model.GpuAnalysis{}, @@ -1535,11 +1562,18 @@ func main() { &model.OperationLog{}, &model.PrequeueConfig{}, &model.QueueQuotaLimit{}, + &model.AgentSession{}, + &model.AgentMessage{}, + &model.AgentToolCall{}, + &model.AgentTurn{}, + &model.AgentRunEvent{}, + &model.JobLogSnapshot{}, + &model.AgentFeedback{}, + &model.OperationLog{}, ) if err != nil { return err } - // create default account account := model.Account{ Name: "default", @@ -1661,6 +1695,15 @@ func main() { } } +func runStatements(tx *gorm.DB, statements []string) error { + for _, stmt := range statements { + if err := tx.Exec(stmt).Error; err != nil { + return err + } + } + return nil +} + type jobListIndex struct { create string drop string diff --git a/backend/cmd/gorm-gen/models/migrate_test.go b/backend/cmd/gorm-gen/models/migrate_test.go deleted file mode 100644 index 70a2eef3d..000000000 --- a/backend/cmd/gorm-gen/models/migrate_test.go +++ /dev/null @@ -1,128 +0,0 @@ -// Copyright 2026 The Crater Project Team, RAIDS-Lab -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package main - -import ( - "testing" - - "gorm.io/driver/sqlite" - "gorm.io/gorm" - - "github.com/raids-lab/crater/dao/model" -) - -//nolint:gocyclo // One integration test verifies forward, repeated, rollback, and repeated rollback behavior. -func TestModelDatasetSourceMigrationAndRollback(t *testing.T) { - db, err := gorm.Open(sqlite.Open("file:model_dataset_source_migration?mode=memory&cache=shared"), &gorm.Config{}) - if err != nil { - t.Fatalf("open sqlite: %v", err) - } - for _, statement := range []string{ - `CREATE TABLE datasets (id integer primary key, name text, url text, type text)`, - `CREATE TABLE model_downloads (id integer primary key, name text, path text, category text)`, - } { - if err := db.Exec(statement).Error; err != nil { - t.Fatalf("create legacy table: %v", err) - } - } - - migration := modelDatasetSourceMigration() - if err := migration.Migrate(db); err != nil { - t.Fatalf("migrate: %v", err) - } - if err := migration.Migrate(db); err != nil { - t.Fatalf("idempotent migrate: %v", err) - } - for _, table := range []any{&model.ModelDatasetSource{}, &model.ModelDatasetDiscovery{}} { - if !db.Migrator().HasTable(table) { - t.Fatalf("missing migrated table %T", table) - } - } - for table, value := range map[string]any{ - "datasets": &model.Dataset{}, - "model_downloads": &model.ModelDownload{}, - } { - if !db.Table(table).Migrator().HasColumn(value, "ModelDatasetSourceID") { - t.Fatalf("%s is missing model_dataset_source_id", table) - } - } - - if err := migration.Rollback(db); err != nil { - t.Fatalf("rollback: %v", err) - } - if err := migration.Rollback(db); err != nil { - t.Fatalf("idempotent rollback: %v", err) - } - if db.Migrator().HasTable(&model.ModelDatasetSource{}) || db.Migrator().HasTable(&model.ModelDatasetDiscovery{}) { - t.Fatal("source or discovery table remains after rollback") - } - if db.Table("datasets").Migrator().HasColumn(&model.Dataset{}, "ModelDatasetSourceID") { - t.Fatal("datasets.model_dataset_source_id remains after rollback") - } - if db.Table("model_downloads").Migrator().HasColumn(&model.ModelDownload{}, "ModelDatasetSourceID") { - t.Fatal("model_downloads.model_dataset_source_id remains after rollback") - } -} - -//nolint:gocyclo // One integration test verifies creation, backfill, idempotency, and rollback. -func TestModelDownloadSubmissionMigrationAndRollback(t *testing.T) { - db, err := gorm.Open(sqlite.Open("file:model_download_submission_migration?mode=memory&cache=shared"), &gorm.Config{}) - if err != nil { - t.Fatalf("open sqlite: %v", err) - } - if err := db.AutoMigrate(&model.ModelDownload{}); err != nil { - t.Fatalf("create model download table: %v", err) - } - active := model.ModelDownload{ - Name: "owner/active", Source: model.ModelSourceModelScope, - Category: model.DownloadCategoryModel, Revision: "main", Path: "public/Models/owner/active", - Status: model.ModelDownloadStatusDownloading, CreatorID: 7, - } - if err := db.Create(&active).Error; err != nil { - t.Fatalf("seed active download: %v", err) - } - migration := modelDownloadSubmissionMigration() - if err := migration.Migrate(db); err != nil { - t.Fatalf("migrate: %v", err) - } - if err := migration.Migrate(db); err != nil { - t.Fatalf("idempotent migrate: %v", err) - } - if !db.Migrator().HasTable(&model.ModelDownloadSubmission{}) { - t.Fatal("model download submission table was not created") - } - for _, field := range []string{"Status", "CompletedAt"} { - if !db.Migrator().HasColumn(&model.ModelDownloadSubmission{}, field) { - t.Fatalf("model download submission table is missing %s", field) - } - } - var reservations []model.ModelDownloadSubmission - if err := db.Where("model_download_id = ?", active.ID).Find(&reservations).Error; err != nil { - t.Fatal(err) - } - if len(reservations) != 1 || reservations[0].UserID != active.CreatorID || - reservations[0].Status != model.ModelDownloadSubmissionReserved { - t.Fatalf("active-download backfill = %+v, want one creator reservation", reservations) - } - if err := migration.Rollback(db); err != nil { - t.Fatalf("rollback: %v", err) - } - if err := migration.Rollback(db); err != nil { - t.Fatalf("idempotent rollback: %v", err) - } - if db.Migrator().HasTable(&model.ModelDownloadSubmission{}) { - t.Fatal("model download submission table remains after rollback") - } -} diff --git a/backend/cmd/model-dataset-governance/main.go b/backend/cmd/model-dataset-governance/main.go deleted file mode 100644 index 4f404542e..000000000 --- a/backend/cmd/model-dataset-governance/main.go +++ /dev/null @@ -1,160 +0,0 @@ -// Copyright 2026 The Crater Project Team, RAIDS-Lab -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package main - -import ( - "context" - "encoding/json" - "flag" - "fmt" - "os" - "os/signal" - "strings" - "syscall" - "time" - - "github.com/raids-lab/crater/dao/query" - "github.com/raids-lab/crater/internal/governance/modeldataset" - "github.com/raids-lab/crater/pkg/config" -) - -var ( - AppVersion = "unknown" - CommitSHA = "unknown" - BuildType = "development" - BuildTime = "unknown" -) - -const ( - defaultStorageRoot = "/crater" - defaultModelsDirectory = "Models" - defaultDatasetsDirectory = "Datasets" - defaultLogicalPrefix = "public" - defaultMaxDepth = 8 - defaultReadmeBytes = 64 * 1024 - defaultScanTimeout = 30 * time.Minute -) - -func main() { - var ( - apply = flag.Bool("apply", false, "write source links and discovery records; default is dry-run") - storageRoot = flag.String("storage-root", envOrDefault("CRATER_STORAGE_ROOT", defaultStorageRoot), "filesystem mount root") - logicalPublicPrefix = flag.String("logical-public-prefix", defaultLogicalPrefix, "logical public prefix used by download records") - physicalPublicPrefix = flag.String("physical-public-prefix", "", "physical public prefix; defaults to storage.prefix.public") - modelsDirectory = flag.String("models-subdirectory", defaultModelsDirectory, "models subdirectory under the public prefix") - modelsDirectories = flag.String( - "models-subdirectories", "", - "optional comma-separated model subdirectories; takes precedence over --models-subdirectory", - ) - datasetsDirectory = flag.String("datasets-subdirectory", defaultDatasetsDirectory, "datasets subdirectory under the public prefix") - maxDepth = flag.Int("max-depth", defaultMaxDepth, "maximum scan depth below each resource root") - excludedDirectories = flag.String( - "exclude-directories", - ".cache,.git,.conda,node_modules,site-packages,test,tests,tmp,temp", - "comma-separated directory basenames to skip", - ) - weightPatterns = flag.String( - "model-weight-patterns", - "*.safetensors,pytorch_model*.bin,model*.bin,*.gguf,tf_model.h5,flax_model.msgpack", - "comma-separated model weight filename patterns", - ) - datasetMarkerPatterns = flag.String( - "dataset-marker-patterns", "", - "comma-separated dataset marker patterns; empty disables filesystem-only dataset discovery", - ) - maxReadmeBytes = flag.Int("max-readme-bytes", defaultReadmeBytes, "maximum bytes read from a local README") - scanTimeout = flag.Duration("scan-timeout", defaultScanTimeout, "maximum filesystem scan duration") - showVersion = flag.Bool("version", false, "print build information and exit") - ) - flag.Parse() - if *showVersion { - fmt.Printf("version=%s commit=%s build_type=%s build_time=%s\n", AppVersion, CommitSHA, BuildType, BuildTime) - return - } - - physicalPrefix := strings.TrimSpace(*physicalPublicPrefix) - if physicalPrefix == "" { - physicalPrefix = config.GetConfig().Storage.Prefix.Public - } - ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM) - defer stop() - ctx, cancel := context.WithTimeout(ctx, *scanTimeout) - defer cancel() - - candidates, err := modeldataset.ScanPublic(ctx, &modeldataset.ScanOptions{ - StorageRoot: *storageRoot, - PublicPrefix: physicalPrefix, - ModelsSubdirectory: *modelsDirectory, - ModelsSubdirectories: splitCSV(*modelsDirectories), - DatasetsSubdirectory: *datasetsDirectory, - MaxDepth: *maxDepth, - ExcludedDirectories: splitCSV(*excludedDirectories), - WeightPatterns: splitCSV(*weightPatterns), - DatasetMarkerPatterns: splitCSV(*datasetMarkerPatterns), - }) - if err != nil { - panic(fmt.Errorf("scan public model and dataset storage: %w", err)) - } - - report, err := modeldataset.ReconcilePublic(ctx, query.GetDB(), candidates, &modeldataset.ReconcileOptions{ - Apply: *apply, - LogicalPublicPrefix: *logicalPublicPrefix, - PhysicalPublicPrefix: physicalPrefix, - PhysicalUserPrefix: config.GetConfig().Storage.Prefix.User, - PhysicalAccountPrefix: config.GetConfig().Storage.Prefix.Account, - MaxReadmeBytes: *maxReadmeBytes, - Now: time.Now(), - }) - if err != nil { - panic(fmt.Errorf("reconcile public model and dataset storage: %w", err)) - } - output := struct { - Mode string `json:"mode"` - StorageRoot string `json:"storageRoot"` - PhysicalPublicPrefix string `json:"physicalPublicPrefix"` - Report modeldataset.ReconcileReport `json:"report"` - }{ - Mode: map[bool]string{true: "apply", false: "dry-run"}[*apply], - StorageRoot: *storageRoot, - PhysicalPublicPrefix: physicalPrefix, - Report: report, - } - encoded, err := json.MarshalIndent(output, "", " ") - if err != nil { - panic(err) - } - fmt.Println(string(encoded)) -} - -func splitCSV(value string) []string { - if strings.TrimSpace(value) == "" { - return nil - } - parts := strings.Split(value, ",") - result := make([]string, 0, len(parts)) - for _, part := range parts { - if part = strings.TrimSpace(part); part != "" { - result = append(result, part) - } - } - return result -} - -func envOrDefault(name, fallback string) string { - if value := strings.TrimSpace(os.Getenv(name)); value != "" { - return value - } - return fallback -} diff --git a/backend/cmd/model-metadata-refresh/main.go b/backend/cmd/model-metadata-refresh/main.go index a22c03621..e8c33072f 100644 --- a/backend/cmd/model-metadata-refresh/main.go +++ b/backend/cmd/model-metadata-refresh/main.go @@ -36,8 +36,6 @@ import ( "github.com/raids-lab/crater/dao/model" "github.com/raids-lab/crater/dao/query" - "github.com/raids-lab/crater/internal/governance/modeldataset" - "github.com/raids-lab/crater/pkg/config" ) const ( @@ -47,8 +45,6 @@ const ( maxSourceDescriptionBytes = 500 maxMetadataTags = 4 maxStoredReadmeBytes = 64 * 1024 - maxModelScopePageBytes = 2 * 1024 * 1024 - maxLogoRedirects = 5 ) type sourceMetadata struct { @@ -71,17 +67,6 @@ type sourceMetadata struct { Tags []string } -type sourceEndpoints struct { - HuggingFace []string - ModelScope []string -} - -type cachedLogo struct { - URL string - Data []byte - ContentType string -} - //nolint:gocyclo // Flag-driven batch orchestration is intentionally kept linear for resumability. func main() { apply := flag.Bool("apply", false, "write refreshed metadata to the database") @@ -89,7 +74,6 @@ func main() { afterID := flag.Uint("after-id", 0, "resume after this model download ID") maxRecords := flag.Int("max-records", 0, "maximum records to process; 0 means unlimited") force := flag.Bool("force", false, "refresh records even if their metadata is still fresh") - missingLogoOnly := flag.Bool("missing-logo-only", false, "refresh only records without a cached source logo") staleAfter := flag.Duration("stale-after", 7*24*time.Hour, "refresh metadata older than this duration") delay := flag.Duration("delay", defaultRequestDelay, "delay between source API requests") source := flag.String("source", "", "optional source filter: huggingface or modelscope") @@ -101,14 +85,9 @@ func main() { panic("source must be huggingface or modelscope") } - appConfig := config.GetConfig() db := query.GetDB() - client := &http.Client{Timeout: time.Duration(appConfig.MetadataTimeoutSeconds()) * time.Second} - endpoints := sourceEndpoints{ - HuggingFace: appConfig.HuggingFaceMetadataEndpoints(), - ModelScope: appConfig.ModelScopeMetadataEndpoints(), - } - avatarCache := make(map[string]cachedLogo) + client := &http.Client{Timeout: 20 * time.Second} + avatarCache := make(map[string]string) var processed, refreshed, failed int cursor := *afterID stop := false @@ -118,23 +97,8 @@ func main() { if *source != "" { builder = builder.Where("source = ?", *source) } - if *missingLogoOnly { - builder = builder.Where(` - model_dataset_source_id IS NULL OR EXISTS ( - SELECT 1 FROM model_dataset_sources - WHERE model_dataset_sources.id = model_downloads.model_dataset_source_id - AND model_dataset_sources.deleted_at IS NULL - AND COALESCE(octet_length(model_dataset_sources.logo_data), 0) = 0 - )`) - } else if !*force { - builder = builder.Where(` - metadata_refreshed_at IS NULL OR metadata_refreshed_at < ? OR - EXISTS ( - SELECT 1 FROM model_dataset_sources - WHERE model_dataset_sources.id = model_downloads.model_dataset_source_id - AND model_dataset_sources.deleted_at IS NULL - AND COALESCE(octet_length(model_dataset_sources.logo_data), 0) = 0 - )`, time.Now().Add(-*staleAfter)) + if !*force { + builder = builder.Where("metadata_refreshed_at IS NULL OR metadata_refreshed_at < ?", time.Now().Add(-*staleAfter)) } if err := builder.Order("id ASC").Limit(*batchSize).Find(&downloads).Error; err != nil { panic(err) @@ -147,10 +111,7 @@ func main() { download := &downloads[i] cursor = download.ID processed++ - if err := refreshDownload( - db, client, endpoints, avatarCache, appConfig.MetadataLogoAllowedHosts(), - appConfig.MetadataMaxLogoBytes(), download, *apply, - ); err != nil { + if err := refreshDownload(db, client, avatarCache, download, *apply); err != nil { failed++ fmt.Printf("FAIL id=%d source=%s name=%s: %v\n", download.ID, download.Source, download.Name, err) } else { @@ -174,48 +135,29 @@ func main() { func refreshDownload( db *gorm.DB, client *http.Client, - endpoints sourceEndpoints, - avatarCache map[string]cachedLogo, - logoAllowedHosts []string, - maxLogoBytes int64, + avatarCache map[string]string, download *model.ModelDownload, apply bool, ) error { - metadata, selectedEndpoint, err := fetchMetadata(client, endpoints, download) + metadata, err := fetchMetadata(client, download) if err != nil { return err } organization := strings.SplitN(download.Name, "/", 2)[0] - organizationKey := strings.ToLower(organization) - logo, ok := avatarCache[organizationKey] - if !ok { - logo, ok = loadCachedOrganizationLogo(db, organization) + logoURL := "" + if download.Source == model.ModelSourceHuggingFace { + var ok bool + logoURL, ok = avatarCache[organization] if !ok { - switch download.Source { - case model.ModelSourceHuggingFace: - logo.URL, err = fetchHuggingFaceAvatar(client, endpoints.HuggingFace, organization) - case model.ModelSourceModelScope: - logo.URL, err = fetchModelScopeAvatar(client, selectedEndpoint, download) - } + logoURL, err = fetchHuggingFaceAvatar(client, organization) if err != nil { - fmt.Printf("WARN id=%d owner avatar lookup failed: %v\n", download.ID, err) - } - if logo.URL != "" { - logo.Data, logo.ContentType, err = fetchLogo( - client, logo.URL, logoAllowedHosts, maxLogoBytes, - ) - if err != nil { - fmt.Printf("WARN id=%d owner avatar cache failed: %v\n", download.ID, err) - // Do not persist an untrusted URL or overwrite a previously cached logo - // when a transient lookup, redirect, or allowlist check fails. - logo = cachedLogo{} - } + return fmt.Errorf("owner avatar lookup: %w", err) } + avatarCache[organization] = logoURL } - avatarCache[organizationKey] = logo } - sourceURL := repositoryURL(download, selectedEndpoint) + sourceURL := repositoryURL(download) fmt.Printf("OK id=%d source=%s name=%s downloads=%d likes=%d tags=%v\n", download.ID, download.Source, download.Name, metadata.Downloads, metadata.Likes, metadata.Tags) if !apply { @@ -224,6 +166,7 @@ func refreshDownload( updates := map[string]any{ "organization": organization, + "logo_url": logoURL, "source_url": sourceURL, "display_name": metadata.DisplayName, "source_description": metadata.Description, @@ -240,9 +183,6 @@ func refreshDownload( "source_likes": metadata.Likes, "metadata_refreshed_at": time.Now(), } - if logo.URL != "" { - updates["logo_url"] = logo.URL - } if metadata.SizeBytes > 0 && download.SizeBytes == 0 { updates["size_bytes"] = metadata.SizeBytes } @@ -252,148 +192,72 @@ func refreshDownload( if metadata.CreatedAt != nil && !metadata.CreatedAt.IsZero() { updates["source_created_at"] = *metadata.CreatedAt } - return db.Transaction(func(tx *gorm.DB) error { - now := time.Now() - sourceRecord := model.ModelDatasetSource{ - Provider: model.ModelDatasetProvider(download.Source), - ResourceType: model.DataType(download.Category), - RepositoryID: download.Name, - RepositoryURL: sourceURL, - Organization: organization, - LogoURL: logo.URL, - LogoData: logo.Data, - LogoContentType: logo.ContentType, - DisplayName: metadata.DisplayName, - Description: metadata.Description, - Readme: metadata.Readme, - License: metadata.License, - Task: metadata.Task, - Library: metadata.Library, - ModelType: metadata.ModelType, - ParameterCount: metadata.ParameterCount, - Private: metadata.Private, - Gated: metadata.Gated, - LoginRequired: metadata.LoginRequired, - Downloads: metadata.Downloads, - Likes: metadata.Likes, - SourceCreatedAt: metadata.CreatedAt, - SourceUpdatedAt: metadata.UpdatedAt, - MetadataRefreshedAt: &now, - } - var persisted model.ModelDatasetSource - lookup := tx.Where( - "provider = ? AND resource_type = ? AND repository_id = ?", - sourceRecord.Provider, sourceRecord.ResourceType, sourceRecord.RepositoryID, - ).First(&persisted) - if errors.Is(lookup.Error, gorm.ErrRecordNotFound) { - if err := tx.Create(&sourceRecord).Error; err != nil { - return fmt.Errorf("create source record: %w", err) - } - persisted = sourceRecord - } else if lookup.Error != nil { - return fmt.Errorf("load source record: %w", lookup.Error) - } else if err := tx.Model(&persisted).Updates(sourceRecord).Error; err != nil { - return fmt.Errorf("update source record: %w", err) - } + if err := db.Model(&model.ModelDownload{}).Where("id = ?", download.ID).Updates(updates).Error; err != nil { + return fmt.Errorf("database update: %w", err) + } - updates["model_dataset_source_id"] = persisted.ID - if err := tx.Model(&model.ModelDownload{}).Where("id = ?", download.ID).Updates(updates).Error; err != nil { - return fmt.Errorf("database update: %w", err) + var dataset model.Dataset + if err := db.Where("name = ? AND type = ?", download.Name, model.DataType(download.Category)).First(&dataset).Error; err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil } + return fmt.Errorf("dataset lookup: %w", err) + } + extra := dataset.Extra.Data() + extra.Tags = mergeTags(extra.Tags, append([]string{string(download.Source)}, metadata.Tags...)) + extra.WebURL = &sourceURL + datasetSize := download.SizeBytes + if datasetSize == 0 { + datasetSize = metadata.SizeBytes + } + datasetUpdates := map[string]any{"extra": datatypes.NewJSONType(extra), "size_bytes": datasetSize} + if metadata.Description != "" && isGeneratedDescription(dataset.Describe, download) { + datasetUpdates["describe"] = metadata.Description + } + if err := db.Model(&model.Dataset{}).Where("id = ?", dataset.ID). + Updates(datasetUpdates).Error; err != nil { + return fmt.Errorf("dataset %d metadata update: %w", dataset.ID, err) + } + return nil +} - physicalPath, public := modeldataset.PhysicalStoragePath( - download.Path, - config.GetConfig().MetadataLogicalPublicPrefix(), - config.GetConfig().Storage.Prefix.Public, - ) - dataset, err := findDatasetForDownload(tx, download, physicalPath, public) +func fetchHuggingFaceAvatar(client *http.Client, owner string) (string, error) { + escapedOwner := url.PathEscape(owner) + endpoints := []string{ + "https://huggingface.co/api/organizations/" + escapedOwner + "/overview", + "https://huggingface.co/api/users/" + escapedOwner + "/overview", + } + for _, endpoint := range endpoints { + response, err := getResponse(client, endpoint) if err != nil { - return err + return "", err } - if dataset == nil { - return nil - } - extra := dataset.Extra.Data() - extra.Tags = mergeTags(extra.Tags, append([]string{string(download.Source)}, metadata.Tags...)) - extra.WebURL = &sourceURL - datasetSize := download.SizeBytes - if datasetSize == 0 { - datasetSize = metadata.SizeBytes - } - datasetUpdates := map[string]any{ - "extra": datatypes.NewJSONType(extra), "size_bytes": datasetSize, - "model_dataset_source_id": persisted.ID, - } - if metadata.Description != "" && isGeneratedDescription(dataset.Describe, download) { - datasetUpdates["describe"] = metadata.Description + if response.StatusCode == http.StatusNotFound { + response.Body.Close() + continue } - if err := tx.Model(&model.Dataset{}).Where("id = ?", dataset.ID).Updates(datasetUpdates).Error; err != nil { - return fmt.Errorf("dataset %d metadata update: %w", dataset.ID, err) + if response.StatusCode != http.StatusOK { + statusCode := response.StatusCode + response.Body.Close() + return "", fmt.Errorf("source returned HTTP %d", statusCode) } - return nil - }) -} - -func findDatasetForDownload( - db *gorm.DB, - download *model.ModelDownload, - physicalPath string, - public bool, -) (*model.Dataset, error) { - if public { - var exact model.Dataset - err := db.Where("url = ? AND type = ?", physicalPath, model.DataType(download.Category)).First(&exact).Error - if err == nil { - return &exact, nil + var payload struct { + AvatarURL string `json:"avatarUrl"` } - if !errors.Is(err, gorm.ErrRecordNotFound) { - return nil, fmt.Errorf("dataset path lookup: %w", err) + decodeErr := json.NewDecoder(response.Body).Decode(&payload) + response.Body.Close() + if decodeErr != nil { + return "", decodeErr } + return payload.AvatarURL, nil } - - var matches []model.Dataset - if err := db.Where( - "name = ? AND type = ? AND model_dataset_source_id IS NULL", - download.Name, model.DataType(download.Category), - ).Order("id ASC").Limit(2).Find(&matches).Error; err != nil { - return nil, fmt.Errorf("dataset identity lookup: %w", err) - } - if len(matches) != 1 { - return nil, nil - } - return &matches[0], nil -} - -func fetchHuggingFaceAvatar(client *http.Client, baseEndpoints []string, owner string) (string, error) { - return modeldataset.FetchHuggingFaceAvatarURL(context.Background(), client, baseEndpoints, owner) + return "", nil } -func fetchModelScopeAvatar( - client *http.Client, baseEndpoint string, download *model.ModelDownload, -) (string, error) { - return modeldataset.FetchModelScopeAvatarURL( - context.Background(), client, repositoryURL(download, baseEndpoint), - ) -} - -func loadCachedOrganizationLogo(db *gorm.DB, organization string) (cachedLogo, bool) { - var sources []model.ModelDatasetSource - err := db.Where( - "LOWER(organization) = ? AND octet_length(logo_data) > 0", strings.ToLower(organization), - ).Order("updated_at DESC").Limit(1).Find(&sources).Error - if err != nil || len(sources) == 0 { - return cachedLogo{}, false - } - source := sources[0] - return cachedLogo{URL: source.LogoURL, Data: source.LogoData, ContentType: source.LogoContentType}, true -} - -func fetchMetadata( - client *http.Client, endpoints sourceEndpoints, download *model.ModelDownload, -) (sourceMetadata, string, error) { +func fetchMetadata(client *http.Client, download *model.ModelDownload) (sourceMetadata, error) { owner, name, ok := strings.Cut(download.Name, "/") if !ok { - return sourceMetadata{}, "", fmt.Errorf("invalid repository name") + return sourceMetadata{}, fmt.Errorf("invalid repository name") } owner = url.PathEscape(owner) name = url.PathEscape(name) @@ -403,6 +267,7 @@ func fetchMetadata( if download.Category == model.DownloadCategoryDataset { resource = "datasets" } + endpoint := fmt.Sprintf("https://huggingface.co/api/%s/%s/%s", resource, owner, name) var payload struct { Downloads int64 `json:"downloads"` Likes int64 `json:"likes"` @@ -422,17 +287,14 @@ func fetchMetadata( Total int64 `json:"total"` } `json:"safetensors"` } - selectedEndpoint, err := getJSONFromEndpoints( - client, endpoints.HuggingFace, "/api/"+resource+"/"+owner+"/"+name, &payload, - ) - if err != nil { - return sourceMetadata{}, "", err + if err := getJSON(client, endpoint, &payload); err != nil { + return sourceMetadata{}, err } var cardData struct { License string `json:"license"` } _ = json.Unmarshal(payload.CardData, &cardData) - readmeURL := selectedEndpoint + "/" + readmeURL := "https://huggingface.co/" if download.Category == model.DownloadCategoryDataset { readmeURL += "datasets/" } @@ -461,13 +323,14 @@ func fetchMetadata( CreatedAt: &payload.CreatedAt, UpdatedAt: &payload.LastModified, Tags: limitTags(payload.Tags), - }, selectedEndpoint, nil + }, nil } resource := "models" if download.Category == model.DownloadCategoryDataset { resource = "datasets" } + endpoint := fmt.Sprintf("https://modelscope.cn/openapi/v1/%s/%s/%s", resource, owner, name) var payload struct { Success bool `json:"success"` Data struct { @@ -488,14 +351,11 @@ func fetchMetadata( LoginRequired bool `json:"login_required"` } `json:"data"` } - selectedEndpoint, err := getJSONFromEndpoints( - client, endpoints.ModelScope, "/openapi/v1/"+resource+"/"+owner+"/"+name, &payload, - ) - if err != nil { - return sourceMetadata{}, "", err + if err := getJSON(client, endpoint, &payload); err != nil { + return sourceMetadata{}, err } if !payload.Success { - return sourceMetadata{}, "", fmt.Errorf("source returned success=false") + return sourceMetadata{}, fmt.Errorf("source returned success=false") } task := "" if len(payload.Data.Tasks) > 0 { @@ -521,7 +381,7 @@ func fetchMetadata( CreatedAt: &payload.Data.CreatedAt, UpdatedAt: &payload.Data.LastModified, Tags: limitTags(append(payload.Data.Tasks, payload.Data.Tags...)), - }, selectedEndpoint, nil + }, nil } func fetchOptionalText(client *http.Client, endpoint string, limit int) (string, error) { @@ -759,44 +619,17 @@ func getJSON(client *http.Client, endpoint string, target any) error { return lastErr } -func getJSONFromEndpoints( - client *http.Client, baseEndpoints []string, path string, target any, -) (string, error) { - var lastErr error - for _, baseEndpoint := range baseEndpoints { - baseEndpoint = strings.TrimRight(baseEndpoint, "/") - if err := getJSON(client, baseEndpoint+path, target); err != nil { - lastErr = err - continue - } - return baseEndpoint, nil - } - if lastErr == nil { - lastErr = errors.New("no source metadata endpoint configured") - } - return "", lastErr -} - -func fetchLogo(client *http.Client, endpoint string, allowedHosts []string, maxBytes int64) ( - data []byte, contentType string, err error, -) { - return modeldataset.FetchSourceLogo( - context.Background(), client, endpoint, allowedHosts, maxBytes, - ) -} - -func repositoryURL(download *model.ModelDownload, baseEndpoint string) string { - baseEndpoint = strings.TrimRight(baseEndpoint, "/") +func repositoryURL(download *model.ModelDownload) string { if download.Source == model.ModelSourceHuggingFace { if download.Category == model.DownloadCategoryDataset { - return baseEndpoint + "/datasets/" + download.Name + return "https://huggingface.co/datasets/" + download.Name } - return baseEndpoint + "/" + download.Name + return "https://huggingface.co/" + download.Name } if download.Category == model.DownloadCategoryDataset { - return baseEndpoint + "/datasets/" + download.Name + return "https://modelscope.cn/datasets/" + download.Name } - return baseEndpoint + "/models/" + download.Name + return "https://modelscope.cn/models/" + download.Name } func limitTags(tags []string) []string { diff --git a/backend/cmd/model-metadata-refresh/main_test.go b/backend/cmd/model-metadata-refresh/main_test.go index dbdbc3a96..5db928722 100644 --- a/backend/cmd/model-metadata-refresh/main_test.go +++ b/backend/cmd/model-metadata-refresh/main_test.go @@ -15,42 +15,11 @@ package main import ( - "bytes" - "encoding/json" - "io" - "net/http" "strings" "testing" "unicode/utf8" - - "gorm.io/driver/sqlite" - "gorm.io/gorm" - - "github.com/raids-lab/crater/dao/model" ) -type roundTripFunc func(*http.Request) (*http.Response, error) - -func (fn roundTripFunc) RoundTrip(request *http.Request) (*http.Response, error) { - return fn(request) -} - -func testHTTPClient(fn roundTripFunc) *http.Client { - return &http.Client{Transport: fn} -} - -func testResponse(status int, contentType string, body []byte) *http.Response { - header := make(http.Header) - if contentType != "" { - header.Set("Content-Type", contentType) - } - return &http.Response{ - StatusCode: status, - Header: header, - Body: io.NopCloser(bytes.NewReader(body)), - } -} - func TestTruncateTextPreservesUTF8(t *testing.T) { text := strings.Repeat("模型介绍", 100) truncated := truncateText(text, 101) @@ -62,162 +31,6 @@ func TestTruncateTextPreservesUTF8(t *testing.T) { } } -func TestGetJSONFromEndpointsFallsBack(t *testing.T) { - payloadBytes, err := json.Marshal(map[string]string{"name": "model"}) - if err != nil { - t.Fatal(err) - } - client := testHTTPClient(func(request *http.Request) (*http.Response, error) { - if request.URL.Host == "unavailable.example" { - return testResponse(http.StatusNotFound, "text/plain", []byte("unavailable")), nil - } - return testResponse(http.StatusOK, "application/json", payloadBytes), nil - }) - var payload struct { - Name string `json:"name"` - } - selected, err := getJSONFromEndpoints( - client, - []string{"https://unavailable.example", "https://working.example"}, - "/api/model", &payload, - ) - if err != nil { - t.Fatalf("getJSONFromEndpoints() error = %v", err) - } - if selected != "https://working.example" || payload.Name != "model" { - t.Fatalf("selected = %q, payload = %#v", selected, payload) - } -} - -func TestFetchLogoCachesOnlyBoundedImages(t *testing.T) { - client := testHTTPClient(func(_ *http.Request) (*http.Response, error) { - return testResponse(http.StatusOK, "image/png", []byte("small-logo")), nil - }) - allowedHosts := []string{"source.example"} - data, contentType, err := fetchLogo(client, "https://source.example/logo.png", allowedHosts, 32) - if err != nil { - t.Fatalf("fetchLogo() error = %v", err) - } - if string(data) != "small-logo" || contentType != "image/png" { - t.Fatalf("data = %q, contentType = %q", data, contentType) - } - if _, _, err := fetchLogo(client, "https://source.example/logo.png", allowedHosts, 4); err == nil { - t.Fatal("fetchLogo() accepted an oversized image") - } -} - -func TestFetchLogoSniffsImageFromGenericContentType(t *testing.T) { - pngHeader := []byte("\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR") - client := testHTTPClient(func(_ *http.Request) (*http.Response, error) { - return testResponse(http.StatusOK, "application/octet-stream", pngHeader), nil - }) - data, contentType, err := fetchLogo( - client, "https://resouces.modelscope.cn/avatar/logo.webp", - []string{"resouces.modelscope.cn"}, 64, - ) - if err != nil { - t.Fatalf("fetchLogo() error = %v", err) - } - if !bytes.Equal(data, pngHeader) || contentType != "image/png" { - t.Fatalf("data = %q, contentType = %q", data, contentType) - } -} - -func TestFetchModelScopeAvatarFromRepositoryPage(t *testing.T) { - client := testHTTPClient(func(request *http.Request) (*http.Response, error) { - if request.URL.Path != "/models/Krea/Krea-2-Turbo" { - t.Fatalf("unexpected path %q", request.URL.Path) - } - body := `` - return testResponse(http.StatusOK, "text/html", []byte(body)), nil - }) - download := &model.ModelDownload{ - Name: "Krea/Krea-2-Turbo", Source: model.ModelSourceModelScope, - Category: model.DownloadCategoryModel, - } - avatarURL, err := fetchModelScopeAvatar(client, "https://modelscope.cn", download) - if err != nil { - t.Fatalf("fetchModelScopeAvatar() error = %v", err) - } - if avatarURL != "https://resources.modelscope.cn/avatar/krea.webp?x-oss-process=image%2Fresize%2Cm_lfit%2Cw_128%2Ch_128" { - t.Fatalf("avatar URL = %q", avatarURL) - } -} - -func TestFetchLogoFollowsOnlyAllowedRedirects(t *testing.T) { - var requestedHosts []string - client := testHTTPClient(func(request *http.Request) (*http.Response, error) { - requestedHosts = append(requestedHosts, request.URL.Hostname()) - if request.URL.Hostname() == "source.example" { - response := testResponse(http.StatusFound, "", nil) - response.Header.Set("Location", "https://cdn.example/logo.png") - return response, nil - } - return testResponse(http.StatusOK, "image/png", []byte("redirected-logo")), nil - }) - - data, _, err := fetchLogo( - client, "https://source.example/logo.png", []string{"source.example", "cdn.example"}, 32, - ) - if err != nil || string(data) != "redirected-logo" { - t.Fatalf("allowed redirect result: data=%q err=%v", data, err) - } - if strings.Join(requestedHosts, ",") != "source.example,cdn.example" { - t.Fatalf("unexpected logo requests: %v", requestedHosts) - } - - requestedHosts = nil - _, _, err = fetchLogo(client, "https://source.example/logo.png", []string{"source.example"}, 32) - if err == nil || strings.Join(requestedHosts, ",") != "source.example" { - t.Fatalf("disallowed redirect was followed: hosts=%v err=%v", requestedHosts, err) - } -} - -func TestFetchLogoRejectsUnsafeInitialURLBeforeRequest(t *testing.T) { - requested := false - client := testHTTPClient(func(_ *http.Request) (*http.Response, error) { - requested = true - return testResponse(http.StatusOK, "image/png", []byte("logo")), nil - }) - - for _, endpoint := range []string{ - "http://cdn-avatars.huggingface.co/logo.png", - "https://169.254.169.254/latest/meta-data", - "https://user:password@cdn-avatars.huggingface.co/logo.png", - } { - if _, _, err := fetchLogo( - client, endpoint, []string{"cdn-avatars.huggingface.co"}, 32, - ); err == nil { - t.Fatalf("fetchLogo() accepted unsafe URL %q", endpoint) - } - } - if requested { - t.Fatal("unsafe logo URL reached the HTTP transport") - } -} - -func TestFindDatasetForDownloadSupportsOnePathlessHistoricalRecord(t *testing.T) { - db, err := gorm.Open(sqlite.Open("file:metadata_pathless?mode=memory&cache=shared"), &gorm.Config{}) - if err != nil { - t.Fatal(err) - } - if err := db.AutoMigrate(&model.ModelDatasetSource{}, &model.Dataset{}); err != nil { - t.Fatal(err) - } - dataset := model.Dataset{Name: "owner/model", URL: "", Type: model.DataTypeModel} - if err := db.Create(&dataset).Error; err != nil { - t.Fatal(err) - } - download := &model.ModelDownload{Name: dataset.Name, Category: model.DownloadCategoryModel} - found, err := findDatasetForDownload(db, download, "shared/Models/owner/model", true) - if err != nil { - t.Fatal(err) - } - if found == nil || found.ID != dataset.ID { - t.Fatalf("findDatasetForDownload() = %#v", found) - } -} - func TestSourceDescriptionStripsMarkupAndPreservesUTF8(t *testing.T) { readme := "

模型介绍

" + strings.Repeat("这是一个中文模型。", 100) description := sourceDescription("", readme) diff --git a/backend/dao/model/agent.go b/backend/dao/model/agent.go new file mode 100644 index 000000000..2daecda51 --- /dev/null +++ b/backend/dao/model/agent.go @@ -0,0 +1,133 @@ +package model + +import ( + "time" + + "gorm.io/datatypes" + "gorm.io/gorm" +) + +// AgentSession represents a conversation session between a user and the agent. +type AgentSession struct { + ID uint `gorm:"primarykey" json:"id"` + SessionID string `gorm:"type:uuid;uniqueIndex;not null" json:"sessionId"` + UserID uint `gorm:"not null" json:"userId"` + AccountID uint `gorm:"not null" json:"accountId"` + Title string `gorm:"type:varchar(255)" json:"title"` + Source string `gorm:"type:varchar(32);not null;default:'chat';index" json:"source"` // chat | ops_audit | system + PageContext datatypes.JSON `json:"pageContext"` + MessageCount int `gorm:"default:0" json:"messageCount"` + LastOrchestrationMode string `gorm:"type:varchar(32);default:'single_agent'" json:"lastOrchestrationMode"` + PinnedAt *time.Time `gorm:"index" json:"pinnedAt,omitempty"` + CreatedAt time.Time `json:"createdAt"` + UpdatedAt time.Time `json:"updatedAt"` + DeletedAt gorm.DeletedAt `gorm:"index" json:"deletedAt"` +} + +// AgentMessage represents an individual message in an agent session. +type AgentMessage struct { + ID uint `gorm:"primarykey" json:"id"` + SessionID string `gorm:"type:uuid;index;not null" json:"sessionId"` + Role string `gorm:"type:varchar(20);not null" json:"role"` // user|assistant|tool + Content string `gorm:"type:text" json:"content"` + ToolCalls datatypes.JSON `json:"toolCalls,omitempty"` + ToolCallID string `gorm:"type:varchar(100)" json:"toolCallId,omitempty"` + ToolName string `gorm:"type:varchar(100)" json:"toolName,omitempty"` + Metadata datatypes.JSON `json:"metadata,omitempty"` + CreatedAt time.Time `json:"createdAt"` +} + +// AgentToolCall is an audit log for tool executions. +type AgentToolCall struct { + ID uint `gorm:"primarykey" json:"id"` + SessionID string `gorm:"type:uuid;index;not null" json:"sessionId"` + TurnID string `gorm:"type:uuid;index" json:"turnId,omitempty"` + MessageID *uint `json:"messageId,omitempty"` + ToolCallID string `gorm:"type:varchar(128);index" json:"toolCallId,omitempty"` + AgentID string `gorm:"type:varchar(128);index" json:"agentId,omitempty"` + ParentEventID *uint `gorm:"index" json:"parentEventId,omitempty"` + AgentRole string `gorm:"type:varchar(32);index" json:"agentRole,omitempty"` + Source string `gorm:"type:varchar(32);not null;default:'backend';index" json:"source"` // backend | local + ToolName string `gorm:"type:varchar(100);not null;index" json:"toolName"` + ToolArgs datatypes.JSON `gorm:"not null" json:"toolArgs"` + ToolResult datatypes.JSON `json:"toolResult,omitempty"` + ResultStatus string `gorm:"type:varchar(32);not null;default:'success'" json:"resultStatus"` + ExecutionBackend string `gorm:"type:varchar(64)" json:"executionBackend,omitempty"` + LatencyMs int `json:"latencyMs,omitempty"` + TokenCount int `json:"tokenCount,omitempty"` + UserConfirmed *bool `json:"userConfirmed,omitempty"` + CreatedAt time.Time `json:"createdAt"` +} + +// AgentTurn represents one agent execution turn within a session. +type AgentTurn struct { + ID uint `gorm:"primarykey" json:"id"` + TurnID string `gorm:"type:uuid;uniqueIndex;not null" json:"turnId"` + SessionID string `gorm:"type:uuid;index;not null" json:"sessionId"` + RequestID string `gorm:"type:varchar(128);index" json:"requestId,omitempty"` + OrchestrationMode string `gorm:"type:varchar(32);default:'single_agent';index" json:"orchestrationMode"` + RootAgentID string `gorm:"type:varchar(128);index" json:"rootAgentId,omitempty"` + Status string `gorm:"type:varchar(32);default:'running';index" json:"status"` + FinalMessageID *uint `gorm:"index" json:"finalMessageId,omitempty"` + Metadata datatypes.JSON `json:"metadata,omitempty"` + StartedAt time.Time `gorm:"index" json:"startedAt"` + EndedAt *time.Time `json:"endedAt,omitempty"` + CreatedAt time.Time `json:"createdAt"` + UpdatedAt time.Time `json:"updatedAt"` +} + +// AgentRunEvent stores semantic events emitted during an agent turn. +type AgentRunEvent struct { + ID uint `gorm:"primarykey" json:"id"` + TurnID string `gorm:"type:uuid;index;not null" json:"turnId"` + SessionID string `gorm:"type:uuid;index;not null" json:"sessionId"` + AgentID string `gorm:"type:varchar(128);index" json:"agentId"` + ParentAgentID string `gorm:"type:varchar(128);index" json:"parentAgentId,omitempty"` + AgentRole string `gorm:"type:varchar(32);index" json:"agentRole"` + EventType string `gorm:"type:varchar(64);index" json:"eventType"` + EventStatus string `gorm:"type:varchar(32);index" json:"eventStatus,omitempty"` + Title string `gorm:"type:varchar(255)" json:"title,omitempty"` + Content string `gorm:"type:text" json:"content,omitempty"` + Metadata datatypes.JSON `json:"metadata,omitempty"` + Sequence int `gorm:"index" json:"sequence"` + StartedAt *time.Time `json:"startedAt,omitempty"` + EndedAt *time.Time `json:"endedAt,omitempty"` + CreatedAt time.Time `json:"createdAt"` +} + +// AgentFeedback stores user feedback (thumbs up/down + optional details) for a message or turn. +type AgentFeedback struct { + ID uint `gorm:"primarykey" json:"id"` + SessionID string `gorm:"type:uuid;index;not null" json:"sessionId"` + UserID uint `gorm:"not null;uniqueIndex:idx_feedback_unique,priority:1" json:"userId"` + AccountID uint `gorm:"not null;index" json:"accountId"` + // message | turn + TargetType string `gorm:"type:varchar(16);not null;uniqueIndex:idx_feedback_unique,priority:2" json:"targetType"` + // message.id or turn_id + TargetID string `gorm:"type:varchar(128);not null;uniqueIndex:idx_feedback_unique,priority:3" json:"targetId"` + // 1 = thumbs up, -1 = thumbs down + Rating int16 `gorm:"not null" json:"rating"` + // ["inaccurate","helpful",...] + Tags datatypes.JSON `json:"tags,omitempty"` + // {"relevance":4,"accuracy":3,...} + Dimensions datatypes.JSON `json:"dimensions,omitempty"` + Comment string `gorm:"type:text" json:"comment,omitempty"` + Status string `gorm:"type:varchar(16);not null;default:'draft'" json:"status"` // draft | submitted + SubmittedAt *time.Time `json:"submittedAt,omitempty"` + EnrichedAt *time.Time `json:"enrichedAt,omitempty"` + CreatedAt time.Time `json:"createdAt"` + UpdatedAt time.Time `json:"updatedAt"` +} + +// JobLogSnapshot is a persisted log snippet captured when a job reaches terminal state. +type JobLogSnapshot struct { + ID uint `gorm:"primarykey" json:"id"` + JobName string `gorm:"type:varchar(255);index;not null" json:"jobName"` + PodName string `gorm:"type:varchar(255);not null" json:"podName"` + ContainerName string `gorm:"type:varchar(255);not null" json:"containerName"` + LogTail string `gorm:"type:text" json:"logTail"` + LogHead string `gorm:"type:text" json:"logHead"` + CapturedAt time.Time `gorm:"not null" json:"capturedAt"` + JobStatus string `gorm:"type:varchar(50)" json:"jobStatus"` + CreatedAt time.Time `json:"createdAt"` +} diff --git a/backend/dao/model/dataset.go b/backend/dao/model/dataset.go index 87cd3f2fd..512659e76 100644 --- a/backend/dao/model/dataset.go +++ b/backend/dao/model/dataset.go @@ -31,9 +31,6 @@ type Dataset struct { MountCount int `gorm:"column:mount_count;not null;default:0;comment:mount count"` SizeBytes int64 `gorm:"not null;default:0;comment:资源文件总大小(字节)"` - ModelDatasetSourceID *uint `gorm:"index;comment:模型或数据集外部来源ID"` - ModelDatasetSource *ModelDatasetSource `gorm:"foreignKey:ModelDatasetSourceID"` - UserDatasets []UserDataset AccountDatasets []AccountDataset } diff --git a/backend/dao/model/modeldatasetsource.go b/backend/dao/model/modeldatasetsource.go deleted file mode 100644 index fe72f82a4..000000000 --- a/backend/dao/model/modeldatasetsource.go +++ /dev/null @@ -1,117 +0,0 @@ -// Copyright 2026 The Crater Project Team, RAIDS-Lab -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package model - -import ( - "time" - - "gorm.io/datatypes" - "gorm.io/gorm" -) - -type ModelDatasetProvider string - -const ( - ModelDatasetProviderHuggingFace ModelDatasetProvider = "huggingface" - ModelDatasetProviderModelScope ModelDatasetProvider = "modelscope" - ModelDatasetProviderExternal ModelDatasetProvider = "external" -) - -// ModelDatasetSource stores upstream repository identity and cached metadata. -// Storage location and sharing remain responsibilities of Dataset. -type ModelDatasetSource struct { - gorm.Model - Provider ModelDatasetProvider `gorm:"type:varchar(64);not null;uniqueIndex:idx_model_dataset_source_identity,priority:1"` - ResourceType DataType `gorm:"type:varchar(32);not null;uniqueIndex:idx_model_dataset_source_identity,priority:2"` - RepositoryID string `gorm:"type:varchar(256);not null;uniqueIndex:idx_model_dataset_source_identity,priority:3"` - RepositoryURL string `gorm:"type:varchar(512);comment:外部仓库页面地址"` - - Organization string `gorm:"type:varchar(128);comment:源站组织或作者"` - LogoURL string `gorm:"type:varchar(512);comment:源站组织头像地址"` - LogoData []byte `gorm:"type:bytea;comment:平台缓存的源站头像"` - LogoContentType string `gorm:"type:varchar(128);comment:平台缓存头像的Content-Type"` - DisplayName string `gorm:"type:varchar(256);comment:源站展示名称"` - Description string `gorm:"type:text;comment:源站简介摘要"` - Readme string `gorm:"type:text;comment:源站README内容(截断保存)"` - License string `gorm:"type:varchar(128);comment:源站许可证"` - Task string `gorm:"type:varchar(128);comment:源站任务分类"` - Library string `gorm:"type:varchar(128);comment:源站框架或库"` - ModelType string `gorm:"type:varchar(128);comment:源站模型类型"` - ParameterCount int64 `gorm:"default:0;comment:模型参数量"` - Private bool `gorm:"default:false;comment:源站是否私有"` - Gated bool `gorm:"default:false;comment:源站是否需要申请访问"` - LoginRequired bool `gorm:"default:false;comment:源站是否要求登录下载"` - Downloads int64 `gorm:"default:0;comment:源站下载次数"` - Likes int64 `gorm:"default:0;comment:源站点赞次数"` - SourceCreatedAt *time.Time `gorm:"comment:源站创建时间"` - SourceUpdatedAt *time.Time `gorm:"comment:源站更新时间"` - MetadataRefreshedAt *time.Time `gorm:"index;comment:源站元数据刷新时间"` -} - -type ModelDatasetDiscoveryStatus string - -const ( - ModelDatasetDiscoveryStatusDiscovered ModelDatasetDiscoveryStatus = "discovered" - ModelDatasetDiscoveryStatusRegistered ModelDatasetDiscoveryStatus = "registered" - ModelDatasetDiscoveryStatusPathMissing ModelDatasetDiscoveryStatus = "path_missing" - ModelDatasetDiscoveryStatusMissing ModelDatasetDiscoveryStatus = "missing" - ModelDatasetDiscoveryStatusIgnored ModelDatasetDiscoveryStatus = "ignored" -) - -type ModelDatasetDiscoveryScope string - -const ( - ModelDatasetDiscoveryScopePublic ModelDatasetDiscoveryScope = "public" - ModelDatasetDiscoveryScopeAccount ModelDatasetDiscoveryScope = "account" - ModelDatasetDiscoveryScopeUser ModelDatasetDiscoveryScope = "user" -) - -type ModelDatasetDiscoveryEvidence struct { - HasConfig bool `json:"hasConfig"` - HasReadme bool `json:"hasReadme"` - WeightFiles int `json:"weightFiles"` - MatchedFiles []string `json:"matchedFiles,omitempty"` - Provider ModelDatasetProvider `json:"provider,omitempty"` - RepositoryID string `json:"repositoryId,omitempty"` - RepositoryURL string `json:"repositoryUrl,omitempty"` - ProvenanceSource string `json:"provenanceSource,omitempty"` - ProvenanceConfidence string `json:"provenanceConfidence,omitempty"` - ConfigNameOrPath string `json:"configNameOrPath,omitempty"` - CandidateURLs []string `json:"candidateUrls,omitempty"` - FilesystemUID string `json:"filesystemUid,omitempty"` - FilesystemGID string `json:"filesystemGid,omitempty"` - ModifiedAt *time.Time `json:"modifiedAt,omitempty"` - OwnerUserID *uint `json:"ownerUserId,omitempty"` - OwnerUsername string `json:"ownerUsername,omitempty"` -} - -// ModelDatasetDiscovery is a non-authoritative filesystem inventory record. -// Discovering a path never grants sharing permissions or publishes a Dataset. -type ModelDatasetDiscovery struct { - gorm.Model - DiscoveryKey string `gorm:"type:varchar(1100);not null;uniqueIndex;comment:稳定发现键"` - Path string `gorm:"type:varchar(1024);index;comment:文件系统路径,允许为空"` - Scope ModelDatasetDiscoveryScope `gorm:"type:varchar(32);not null;index"` - ScopeID *uint `gorm:"index;comment:用户或队列ID"` - DetectedType DataType `gorm:"type:varchar(32);not null;index"` - DetectedName string `gorm:"type:varchar(256);not null"` - Evidence datatypes.JSONType[ModelDatasetDiscoveryEvidence] `gorm:"comment:文件系统检测依据"` - SizeBytes int64 `gorm:"not null;default:0"` - DatasetID *uint `gorm:"index"` - SourceID *uint `gorm:"index"` - Status ModelDatasetDiscoveryStatus `gorm:"type:varchar(32);not null;default:discovered;index"` - FirstSeenAt time.Time `gorm:"not null"` - LastSeenAt time.Time `gorm:"not null;index"` -} diff --git a/backend/dao/model/modeldownload.go b/backend/dao/model/modeldownload.go index 0337994cd..7f776c874 100644 --- a/backend/dao/model/modeldownload.go +++ b/backend/dao/model/modeldownload.go @@ -37,43 +37,41 @@ type ModelDownload struct { Source ModelSource `gorm:"type:varchar(32);not null;default:modelscope;uniqueIndex:idx_download_unique,priority:2"` Category DownloadCategory `gorm:"type:varchar(32);not null;default:model;uniqueIndex:idx_download_unique,priority:3"` - Revision string `gorm:"type:varchar(128);uniqueIndex:idx_download_unique,priority:4;comment:版本/分支/commit"` - Path string `gorm:"type:varchar(512);not null;comment:实际下载路径"` - SizeBytes int64 `gorm:"default:0;comment:文件总大小(字节)"` - DownloadedBytes int64 `gorm:"default:0;comment:已下载大小(字节)"` - DownloadSpeed string `gorm:"type:varchar(32);comment:下载速度(如: 10MB/s)"` - Status ModelDownloadStatus `gorm:"type:varchar(32);not null;default:Pending;comment:下载状态"` - Message string `gorm:"type:text;comment:状态消息(错误信息等)"` - Logs string `gorm:"type:text;comment:终态时保存的Pod日志(K8s Job被GC后仍可查看)"` - LogsSavedAt *time.Time `gorm:"comment:日志保存时间"` - Organization string `gorm:"type:varchar(128);comment:源站组织或作者"` - LogoURL string `gorm:"type:varchar(512);comment:源站组织头像地址"` - SourceURL string `gorm:"type:varchar(512);comment:源站仓库详情地址"` - DisplayName string `gorm:"type:varchar(256);comment:源站展示名称"` - SourceDescription string `gorm:"type:text;comment:源站简介摘要"` - SourceReadme string `gorm:"type:text;comment:源站README内容(截断保存)"` - License string `gorm:"type:varchar(128);comment:源站许可证"` - Task string `gorm:"type:varchar(128);comment:源站任务分类"` - Library string `gorm:"type:varchar(128);comment:源站框架或库"` - ModelType string `gorm:"type:varchar(128);comment:源站模型类型"` - ParameterCount int64 `gorm:"default:0;comment:模型参数量"` - SourcePrivate bool `gorm:"default:false;comment:源站是否私有"` - SourceGated bool `gorm:"default:false;comment:源站是否需要申请访问"` - SourceLoginRequired bool `gorm:"default:false;comment:源站是否要求登录下载"` - SourceDownloads int64 `gorm:"default:0;comment:源站下载次数"` - SourceLikes int64 `gorm:"default:0;comment:源站点赞次数"` - SourceCreatedAt *time.Time `gorm:"comment:源站创建时间"` - SourceUpdatedAt *time.Time `gorm:"comment:源站更新时间"` - MetadataRefreshedAt *time.Time `gorm:"comment:源站元数据刷新时间"` - ModelDatasetSourceID *uint `gorm:"index;comment:模型或数据集外部来源ID"` - ModelDatasetSource *ModelDatasetSource `gorm:"foreignKey:ModelDatasetSourceID"` - JobName string `gorm:"type:varchar(256);comment:K8s Job名称"` - CreatorID uint `gorm:"not null;comment:首个发起下载的用户ID"` - Creator User `gorm:"foreignKey:CreatorID"` - ReferenceCount int `gorm:"default:0;comment:提交下载需求的用户计数"` + Revision string `gorm:"type:varchar(128);uniqueIndex:idx_download_unique,priority:4;comment:版本/分支/commit"` + Path string `gorm:"type:varchar(512);not null;comment:实际下载路径"` + SizeBytes int64 `gorm:"default:0;comment:文件总大小(字节)"` + DownloadedBytes int64 `gorm:"default:0;comment:已下载大小(字节)"` + DownloadSpeed string `gorm:"type:varchar(32);comment:下载速度(如: 10MB/s)"` + Status ModelDownloadStatus `gorm:"type:varchar(32);not null;default:Pending;comment:下载状态"` + Message string `gorm:"type:text;comment:状态消息(错误信息等)"` + Logs string `gorm:"type:text;comment:终态时保存的Pod日志(K8s Job被GC后仍可查看)"` + LogsSavedAt *time.Time `gorm:"comment:日志保存时间"` + Organization string `gorm:"type:varchar(128);comment:源站组织或作者"` + LogoURL string `gorm:"type:varchar(512);comment:源站组织头像地址"` + SourceURL string `gorm:"type:varchar(512);comment:源站仓库详情地址"` + DisplayName string `gorm:"type:varchar(256);comment:源站展示名称"` + SourceDescription string `gorm:"type:text;comment:源站简介摘要"` + SourceReadme string `gorm:"type:text;comment:源站README内容(截断保存)"` + License string `gorm:"type:varchar(128);comment:源站许可证"` + Task string `gorm:"type:varchar(128);comment:源站任务分类"` + Library string `gorm:"type:varchar(128);comment:源站框架或库"` + ModelType string `gorm:"type:varchar(128);comment:源站模型类型"` + ParameterCount int64 `gorm:"default:0;comment:模型参数量"` + SourcePrivate bool `gorm:"default:false;comment:源站是否私有"` + SourceGated bool `gorm:"default:false;comment:源站是否需要申请访问"` + SourceLoginRequired bool `gorm:"default:false;comment:源站是否要求登录下载"` + SourceDownloads int64 `gorm:"default:0;comment:源站下载次数"` + SourceLikes int64 `gorm:"default:0;comment:源站点赞次数"` + SourceCreatedAt *time.Time `gorm:"comment:源站创建时间"` + SourceUpdatedAt *time.Time `gorm:"comment:源站更新时间"` + MetadataRefreshedAt *time.Time `gorm:"comment:源站元数据刷新时间"` + JobName string `gorm:"type:varchar(256);comment:K8s Job名称"` + CreatorID uint `gorm:"not null;comment:首个发起下载的用户ID"` + Creator User `gorm:"foreignKey:CreatorID"` + ReferenceCount int `gorm:"default:0;comment:引用计数"` } -// UserModelDownload records users who explicitly submitted a need for this download. +// UserModelDownload 用户与下载的关联表 type UserModelDownload struct { ID uint `gorm:"primaryKey"` UserID uint `gorm:"not null;comment:用户ID;uniqueIndex:idx_user_download,priority:1"` @@ -82,38 +80,3 @@ type UserModelDownload struct { User User `gorm:"foreignKey:UserID"` ModelDownload ModelDownload `gorm:"foreignKey:ModelDownloadID"` } - -type ModelDownloadSubmissionAction string - -const ( - ModelDownloadSubmissionCreate ModelDownloadSubmissionAction = "create" - ModelDownloadSubmissionRetry ModelDownloadSubmissionAction = "retry" - ModelDownloadSubmissionResume ModelDownloadSubmissionAction = "resume" -) - -type ModelDownloadSubmissionStatus string - -const ( - // ModelDownloadSubmissionReserved temporarily occupies one rolling-window - // slot while the Kubernetes download Job is active. - ModelDownloadSubmissionReserved ModelDownloadSubmissionStatus = "Reserved" - // ModelDownloadSubmissionSucceeded starts consuming the rolling window from - // the time the model or dataset finishes downloading. - ModelDownloadSubmissionSucceeded ModelDownloadSubmissionStatus = "Succeeded" - // ModelDownloadSubmissionReleased records an attempt that failed, was paused, - // or was otherwise canceled and therefore does not consume quota. - ModelDownloadSubmissionReleased ModelDownloadSubmissionStatus = "Released" -) - -// ModelDownloadSubmission records quota reservations for Jobs that may produce -// a completed model or dataset. Reusing an existing public download does not -// create a quota submission. -type ModelDownloadSubmission struct { - ID uint `gorm:"primaryKey"` - UserID uint `gorm:"not null;index:idx_mds_uc,priority:1;index:idx_mds_q,priority:1"` - ModelDownloadID uint `gorm:"not null;index"` - Action ModelDownloadSubmissionAction `gorm:"type:varchar(16);not null"` - Status ModelDownloadSubmissionStatus `gorm:"type:varchar(16);not null;default:Reserved;index:idx_mds_q,priority:2"` - CreatedAt time.Time `gorm:"not null;index:idx_mds_uc,priority:2"` - CompletedAt *time.Time `gorm:"index:idx_mds_q,priority:3"` -} diff --git a/backend/dao/model/system_config.go b/backend/dao/model/system_config.go index 92df85d72..4da4b8006 100644 --- a/backend/dao/model/system_config.go +++ b/backend/dao/model/system_config.go @@ -27,13 +27,6 @@ const ( ConfigKeyBillingDefaultIssuePeriodMinute = "BILLING_DEFAULT_ISSUE_PERIOD_MINUTES" ConfigKeyBillingAccountIssueAmountOverrideEnabled = "ENABLE_BILLING_ACCOUNT_ISSUE_AMOUNT_OVERRIDE" ConfigKeyBillingAccountIssuePeriodOverrideEnabled = "ENABLE_BILLING_ACCOUNT_ISSUE_PERIOD_OVERRIDE" - - // 模型与数据集下载额度配置键 - ConfigKeyModelDownloadLimitEnabled = "MODEL_DOWNLOAD_LIMIT_ENABLED" - ConfigKeyModelDownloadMaxConcurrent = "MODEL_DOWNLOAD_MAX_CONCURRENT" - ConfigKeyModelDownloadWindowHours = "MODEL_DOWNLOAD_WINDOW_HOURS" - ConfigKeyModelDownloadMaxSuccessfulDownloads = "MODEL_DOWNLOAD_MAX_SUCCESSFUL_DOWNLOADS" - ConfigKeyModelDownloadWhitelistUsers = "MODEL_DOWNLOAD_WHITELIST_USER_IDS" ) // DefaultConfigKeys 定义了系统启动时必须存在的键 @@ -51,9 +44,4 @@ var DefaultConfigKeys = []string{ ConfigKeyBillingDefaultIssuePeriodMinute, ConfigKeyBillingAccountIssueAmountOverrideEnabled, ConfigKeyBillingAccountIssuePeriodOverrideEnabled, - ConfigKeyModelDownloadLimitEnabled, - ConfigKeyModelDownloadMaxConcurrent, - ConfigKeyModelDownloadWindowHours, - ConfigKeyModelDownloadMaxSuccessfulDownloads, - ConfigKeyModelDownloadWhitelistUsers, } diff --git a/backend/dao/query/accounts.gen.go b/backend/dao/query/accounts.gen.go index f39288892..22cb2b21c 100644 --- a/backend/dao/query/accounts.gen.go +++ b/backend/dao/query/accounts.gen.go @@ -72,7 +72,7 @@ type account struct { ExpiredAt field.Time // 账户过期时间 Quota field.Field // 账户对应队列的资源配额 UserDefaultQuota field.Field // 账户中用户默认的资源配额模版 - BillingIssueAmount field.Int64 // 账户周期发放点数额度(内部微点, 为空表示未配置) + BillingIssueAmount field.Int64 // 账户周期发放点数额度(为空表示未配置) BillingIssuePeriodMinutes field.Int // 账户周期发放间隔分钟(<=0表示关闭, 为空表示未配置) BillingLastIssuedAt field.Time // 账户上次发放时间 UserAccounts accountHasManyUserAccounts diff --git a/backend/dao/query/datasets.gen.go b/backend/dao/query/datasets.gen.go index ca979482d..afbc2daac 100644 --- a/backend/dao/query/datasets.gen.go +++ b/backend/dao/query/datasets.gen.go @@ -40,7 +40,6 @@ func newDataset(db *gorm.DB, opts ...gen.DOOption) dataset { _dataset.UserID = field.NewUint(tableName, "user_id") _dataset.MountCount = field.NewInt(tableName, "mount_count") _dataset.SizeBytes = field.NewInt64(tableName, "size_bytes") - _dataset.ModelDatasetSourceID = field.NewUint(tableName, "model_dataset_source_id") _dataset.UserDatasets = datasetHasManyUserDatasets{ db: db.Session(&gorm.Session{}), @@ -69,12 +68,6 @@ func newDataset(db *gorm.DB, opts ...gen.DOOption) dataset { }, } - _dataset.ModelDatasetSource = datasetBelongsToModelDatasetSource{ - db: db.Session(&gorm.Session{}), - - RelationField: field.NewRelation("ModelDatasetSource", "model.ModelDatasetSource"), - } - _dataset.fillFieldMap() return _dataset @@ -83,28 +76,25 @@ func newDataset(db *gorm.DB, opts ...gen.DOOption) dataset { type dataset struct { datasetDo datasetDo - ALL field.Asterisk - ID field.Uint - CreatedAt field.Time - UpdatedAt field.Time - DeletedAt field.Field - Name field.String // 数据集名 - URL field.String // 数据集空间路径 - Describe field.String // 数据集描述 - Type field.String // 数据类型 - Extra field.Field // 额外信息(tags、weburl等) - UserID field.Uint - MountCount field.Int // mount count - SizeBytes field.Int64 // 资源文件总大小(字节) - ModelDatasetSourceID field.Uint // 模型或数据集外部来源ID - UserDatasets datasetHasManyUserDatasets + ALL field.Asterisk + ID field.Uint + CreatedAt field.Time + UpdatedAt field.Time + DeletedAt field.Field + Name field.String // 数据集名 + URL field.String // 数据集空间路径 + Describe field.String // 数据集描述 + Type field.String // 数据类型 + Extra field.Field // 额外信息(tags、weburl等) + UserID field.Uint + MountCount field.Int // mount count + SizeBytes field.Int64 // 资源文件总大小(字节) + UserDatasets datasetHasManyUserDatasets AccountDatasets datasetHasManyAccountDatasets User datasetBelongsToUser - ModelDatasetSource datasetBelongsToModelDatasetSource - fieldMap map[string]field.Expr } @@ -132,7 +122,6 @@ func (d *dataset) updateTableName(table string) *dataset { d.UserID = field.NewUint(table, "user_id") d.MountCount = field.NewInt(table, "mount_count") d.SizeBytes = field.NewInt64(table, "size_bytes") - d.ModelDatasetSourceID = field.NewUint(table, "model_dataset_source_id") d.fillFieldMap() @@ -157,7 +146,7 @@ func (d *dataset) GetFieldByName(fieldName string) (field.OrderExpr, bool) { } func (d *dataset) fillFieldMap() { - d.fieldMap = make(map[string]field.Expr, 17) + d.fieldMap = make(map[string]field.Expr, 15) d.fieldMap["id"] = d.ID d.fieldMap["created_at"] = d.CreatedAt d.fieldMap["updated_at"] = d.UpdatedAt @@ -170,7 +159,6 @@ func (d *dataset) fillFieldMap() { d.fieldMap["user_id"] = d.UserID d.fieldMap["mount_count"] = d.MountCount d.fieldMap["size_bytes"] = d.SizeBytes - d.fieldMap["model_dataset_source_id"] = d.ModelDatasetSourceID } @@ -182,8 +170,6 @@ func (d dataset) clone(db *gorm.DB) dataset { d.AccountDatasets.db.Statement.ConnPool = db.Statement.ConnPool d.User.db = db.Session(&gorm.Session{Initialized: true}) d.User.db.Statement.ConnPool = db.Statement.ConnPool - d.ModelDatasetSource.db = db.Session(&gorm.Session{Initialized: true}) - d.ModelDatasetSource.db.Statement.ConnPool = db.Statement.ConnPool return d } @@ -192,7 +178,6 @@ func (d dataset) replaceDB(db *gorm.DB) dataset { d.UserDatasets.db = db.Session(&gorm.Session{}) d.AccountDatasets.db = db.Session(&gorm.Session{}) d.User.db = db.Session(&gorm.Session{}) - d.ModelDatasetSource.db = db.Session(&gorm.Session{}) return d } @@ -446,87 +431,6 @@ func (a datasetBelongsToUserTx) Unscoped() *datasetBelongsToUserTx { return &a } -type datasetBelongsToModelDatasetSource struct { - db *gorm.DB - - field.RelationField -} - -func (a datasetBelongsToModelDatasetSource) Where(conds ...field.Expr) *datasetBelongsToModelDatasetSource { - if len(conds) == 0 { - return &a - } - - exprs := make([]clause.Expression, 0, len(conds)) - for _, cond := range conds { - exprs = append(exprs, cond.BeCond().(clause.Expression)) - } - a.db = a.db.Clauses(clause.Where{Exprs: exprs}) - return &a -} - -func (a datasetBelongsToModelDatasetSource) WithContext(ctx context.Context) *datasetBelongsToModelDatasetSource { - a.db = a.db.WithContext(ctx) - return &a -} - -func (a datasetBelongsToModelDatasetSource) Session(session *gorm.Session) *datasetBelongsToModelDatasetSource { - a.db = a.db.Session(session) - return &a -} - -func (a datasetBelongsToModelDatasetSource) Model(m *model.Dataset) *datasetBelongsToModelDatasetSourceTx { - return &datasetBelongsToModelDatasetSourceTx{a.db.Model(m).Association(a.Name())} -} - -func (a datasetBelongsToModelDatasetSource) Unscoped() *datasetBelongsToModelDatasetSource { - a.db = a.db.Unscoped() - return &a -} - -type datasetBelongsToModelDatasetSourceTx struct{ tx *gorm.Association } - -func (a datasetBelongsToModelDatasetSourceTx) Find() (result *model.ModelDatasetSource, err error) { - return result, a.tx.Find(&result) -} - -func (a datasetBelongsToModelDatasetSourceTx) Append(values ...*model.ModelDatasetSource) (err error) { - targetValues := make([]interface{}, len(values)) - for i, v := range values { - targetValues[i] = v - } - return a.tx.Append(targetValues...) -} - -func (a datasetBelongsToModelDatasetSourceTx) Replace(values ...*model.ModelDatasetSource) (err error) { - targetValues := make([]interface{}, len(values)) - for i, v := range values { - targetValues[i] = v - } - return a.tx.Replace(targetValues...) -} - -func (a datasetBelongsToModelDatasetSourceTx) Delete(values ...*model.ModelDatasetSource) (err error) { - targetValues := make([]interface{}, len(values)) - for i, v := range values { - targetValues[i] = v - } - return a.tx.Delete(targetValues...) -} - -func (a datasetBelongsToModelDatasetSourceTx) Clear() error { - return a.tx.Clear() -} - -func (a datasetBelongsToModelDatasetSourceTx) Count() int64 { - return a.tx.Count() -} - -func (a datasetBelongsToModelDatasetSourceTx) Unscoped() *datasetBelongsToModelDatasetSourceTx { - a.tx = a.tx.Unscoped() - return &a -} - type datasetDo struct{ gen.DO } type IDatasetDo interface { diff --git a/backend/dao/query/gen.go b/backend/dao/query/gen.go index 349ef1fd4..fccb85359 100644 --- a/backend/dao/query/gen.go +++ b/backend/dao/query/gen.go @@ -16,37 +16,34 @@ import ( ) var ( - Q = new(Query) - AITask *aITask - Account *account - AccountDataset *accountDataset - Alert *alert - ApprovalOrder *approvalOrder - CronJobConfig *cronJobConfig - CronJobRecord *cronJobRecord - CudaBaseImage *cudaBaseImage - Dataset *dataset - GpuAnalysis *gpuAnalysis - Image *image - ImageAccount *imageAccount - ImageUser *imageUser - Job *job - Jobtemplate *jobtemplate - Kaniko *kaniko - ModelDatasetDiscovery *modelDatasetDiscovery - ModelDatasetSource *modelDatasetSource - ModelDownload *modelDownload - ModelDownloadSubmission *modelDownloadSubmission - PrequeueConfig *prequeueConfig - QueueQuotaLimit *queueQuotaLimit - Resource *resource - ResourceNetwork *resourceNetwork - ResourceVGPU *resourceVGPU - SystemConfig *systemConfig - User *user - UserAccount *userAccount - UserDataset *userDataset - UserModelDownload *userModelDownload + Q = new(Query) + AITask *aITask + Account *account + AccountDataset *accountDataset + Alert *alert + ApprovalOrder *approvalOrder + CronJobConfig *cronJobConfig + CronJobRecord *cronJobRecord + CudaBaseImage *cudaBaseImage + Dataset *dataset + GpuAnalysis *gpuAnalysis + Image *image + ImageAccount *imageAccount + ImageUser *imageUser + Job *job + Jobtemplate *jobtemplate + Kaniko *kaniko + ModelDownload *modelDownload + PrequeueConfig *prequeueConfig + QueueQuotaLimit *queueQuotaLimit + Resource *resource + ResourceNetwork *resourceNetwork + ResourceVGPU *resourceVGPU + SystemConfig *systemConfig + User *user + UserAccount *userAccount + UserDataset *userDataset + UserModelDownload *userModelDownload ) func SetDefault(db *gorm.DB, opts ...gen.DOOption) { @@ -67,10 +64,7 @@ func SetDefault(db *gorm.DB, opts ...gen.DOOption) { Job = &Q.Job Jobtemplate = &Q.Jobtemplate Kaniko = &Q.Kaniko - ModelDatasetDiscovery = &Q.ModelDatasetDiscovery - ModelDatasetSource = &Q.ModelDatasetSource ModelDownload = &Q.ModelDownload - ModelDownloadSubmission = &Q.ModelDownloadSubmission PrequeueConfig = &Q.PrequeueConfig QueueQuotaLimit = &Q.QueueQuotaLimit Resource = &Q.Resource @@ -85,110 +79,101 @@ func SetDefault(db *gorm.DB, opts ...gen.DOOption) { func Use(db *gorm.DB, opts ...gen.DOOption) *Query { return &Query{ - db: db, - AITask: newAITask(db, opts...), - Account: newAccount(db, opts...), - AccountDataset: newAccountDataset(db, opts...), - Alert: newAlert(db, opts...), - ApprovalOrder: newApprovalOrder(db, opts...), - CronJobConfig: newCronJobConfig(db, opts...), - CronJobRecord: newCronJobRecord(db, opts...), - CudaBaseImage: newCudaBaseImage(db, opts...), - Dataset: newDataset(db, opts...), - GpuAnalysis: newGpuAnalysis(db, opts...), - Image: newImage(db, opts...), - ImageAccount: newImageAccount(db, opts...), - ImageUser: newImageUser(db, opts...), - Job: newJob(db, opts...), - Jobtemplate: newJobtemplate(db, opts...), - Kaniko: newKaniko(db, opts...), - ModelDatasetDiscovery: newModelDatasetDiscovery(db, opts...), - ModelDatasetSource: newModelDatasetSource(db, opts...), - ModelDownload: newModelDownload(db, opts...), - ModelDownloadSubmission: newModelDownloadSubmission(db, opts...), - PrequeueConfig: newPrequeueConfig(db, opts...), - QueueQuotaLimit: newQueueQuotaLimit(db, opts...), - Resource: newResource(db, opts...), - ResourceNetwork: newResourceNetwork(db, opts...), - ResourceVGPU: newResourceVGPU(db, opts...), - SystemConfig: newSystemConfig(db, opts...), - User: newUser(db, opts...), - UserAccount: newUserAccount(db, opts...), - UserDataset: newUserDataset(db, opts...), - UserModelDownload: newUserModelDownload(db, opts...), + db: db, + AITask: newAITask(db, opts...), + Account: newAccount(db, opts...), + AccountDataset: newAccountDataset(db, opts...), + Alert: newAlert(db, opts...), + ApprovalOrder: newApprovalOrder(db, opts...), + CronJobConfig: newCronJobConfig(db, opts...), + CronJobRecord: newCronJobRecord(db, opts...), + CudaBaseImage: newCudaBaseImage(db, opts...), + Dataset: newDataset(db, opts...), + GpuAnalysis: newGpuAnalysis(db, opts...), + Image: newImage(db, opts...), + ImageAccount: newImageAccount(db, opts...), + ImageUser: newImageUser(db, opts...), + Job: newJob(db, opts...), + Jobtemplate: newJobtemplate(db, opts...), + Kaniko: newKaniko(db, opts...), + ModelDownload: newModelDownload(db, opts...), + PrequeueConfig: newPrequeueConfig(db, opts...), + QueueQuotaLimit: newQueueQuotaLimit(db, opts...), + Resource: newResource(db, opts...), + ResourceNetwork: newResourceNetwork(db, opts...), + ResourceVGPU: newResourceVGPU(db, opts...), + SystemConfig: newSystemConfig(db, opts...), + User: newUser(db, opts...), + UserAccount: newUserAccount(db, opts...), + UserDataset: newUserDataset(db, opts...), + UserModelDownload: newUserModelDownload(db, opts...), } } type Query struct { db *gorm.DB - AITask aITask - Account account - AccountDataset accountDataset - Alert alert - ApprovalOrder approvalOrder - CronJobConfig cronJobConfig - CronJobRecord cronJobRecord - CudaBaseImage cudaBaseImage - Dataset dataset - GpuAnalysis gpuAnalysis - Image image - ImageAccount imageAccount - ImageUser imageUser - Job job - Jobtemplate jobtemplate - Kaniko kaniko - ModelDatasetDiscovery modelDatasetDiscovery - ModelDatasetSource modelDatasetSource - ModelDownload modelDownload - ModelDownloadSubmission modelDownloadSubmission - PrequeueConfig prequeueConfig - QueueQuotaLimit queueQuotaLimit - Resource resource - ResourceNetwork resourceNetwork - ResourceVGPU resourceVGPU - SystemConfig systemConfig - User user - UserAccount userAccount - UserDataset userDataset - UserModelDownload userModelDownload + AITask aITask + Account account + AccountDataset accountDataset + Alert alert + ApprovalOrder approvalOrder + CronJobConfig cronJobConfig + CronJobRecord cronJobRecord + CudaBaseImage cudaBaseImage + Dataset dataset + GpuAnalysis gpuAnalysis + Image image + ImageAccount imageAccount + ImageUser imageUser + Job job + Jobtemplate jobtemplate + Kaniko kaniko + ModelDownload modelDownload + PrequeueConfig prequeueConfig + QueueQuotaLimit queueQuotaLimit + Resource resource + ResourceNetwork resourceNetwork + ResourceVGPU resourceVGPU + SystemConfig systemConfig + User user + UserAccount userAccount + UserDataset userDataset + UserModelDownload userModelDownload } func (q *Query) Available() bool { return q.db != nil } func (q *Query) clone(db *gorm.DB) *Query { return &Query{ - db: db, - AITask: q.AITask.clone(db), - Account: q.Account.clone(db), - AccountDataset: q.AccountDataset.clone(db), - Alert: q.Alert.clone(db), - ApprovalOrder: q.ApprovalOrder.clone(db), - CronJobConfig: q.CronJobConfig.clone(db), - CronJobRecord: q.CronJobRecord.clone(db), - CudaBaseImage: q.CudaBaseImage.clone(db), - Dataset: q.Dataset.clone(db), - GpuAnalysis: q.GpuAnalysis.clone(db), - Image: q.Image.clone(db), - ImageAccount: q.ImageAccount.clone(db), - ImageUser: q.ImageUser.clone(db), - Job: q.Job.clone(db), - Jobtemplate: q.Jobtemplate.clone(db), - Kaniko: q.Kaniko.clone(db), - ModelDatasetDiscovery: q.ModelDatasetDiscovery.clone(db), - ModelDatasetSource: q.ModelDatasetSource.clone(db), - ModelDownload: q.ModelDownload.clone(db), - ModelDownloadSubmission: q.ModelDownloadSubmission.clone(db), - PrequeueConfig: q.PrequeueConfig.clone(db), - QueueQuotaLimit: q.QueueQuotaLimit.clone(db), - Resource: q.Resource.clone(db), - ResourceNetwork: q.ResourceNetwork.clone(db), - ResourceVGPU: q.ResourceVGPU.clone(db), - SystemConfig: q.SystemConfig.clone(db), - User: q.User.clone(db), - UserAccount: q.UserAccount.clone(db), - UserDataset: q.UserDataset.clone(db), - UserModelDownload: q.UserModelDownload.clone(db), + db: db, + AITask: q.AITask.clone(db), + Account: q.Account.clone(db), + AccountDataset: q.AccountDataset.clone(db), + Alert: q.Alert.clone(db), + ApprovalOrder: q.ApprovalOrder.clone(db), + CronJobConfig: q.CronJobConfig.clone(db), + CronJobRecord: q.CronJobRecord.clone(db), + CudaBaseImage: q.CudaBaseImage.clone(db), + Dataset: q.Dataset.clone(db), + GpuAnalysis: q.GpuAnalysis.clone(db), + Image: q.Image.clone(db), + ImageAccount: q.ImageAccount.clone(db), + ImageUser: q.ImageUser.clone(db), + Job: q.Job.clone(db), + Jobtemplate: q.Jobtemplate.clone(db), + Kaniko: q.Kaniko.clone(db), + ModelDownload: q.ModelDownload.clone(db), + PrequeueConfig: q.PrequeueConfig.clone(db), + QueueQuotaLimit: q.QueueQuotaLimit.clone(db), + Resource: q.Resource.clone(db), + ResourceNetwork: q.ResourceNetwork.clone(db), + ResourceVGPU: q.ResourceVGPU.clone(db), + SystemConfig: q.SystemConfig.clone(db), + User: q.User.clone(db), + UserAccount: q.UserAccount.clone(db), + UserDataset: q.UserDataset.clone(db), + UserModelDownload: q.UserModelDownload.clone(db), } } @@ -202,105 +187,96 @@ func (q *Query) WriteDB() *Query { func (q *Query) ReplaceDB(db *gorm.DB) *Query { return &Query{ - db: db, - AITask: q.AITask.replaceDB(db), - Account: q.Account.replaceDB(db), - AccountDataset: q.AccountDataset.replaceDB(db), - Alert: q.Alert.replaceDB(db), - ApprovalOrder: q.ApprovalOrder.replaceDB(db), - CronJobConfig: q.CronJobConfig.replaceDB(db), - CronJobRecord: q.CronJobRecord.replaceDB(db), - CudaBaseImage: q.CudaBaseImage.replaceDB(db), - Dataset: q.Dataset.replaceDB(db), - GpuAnalysis: q.GpuAnalysis.replaceDB(db), - Image: q.Image.replaceDB(db), - ImageAccount: q.ImageAccount.replaceDB(db), - ImageUser: q.ImageUser.replaceDB(db), - Job: q.Job.replaceDB(db), - Jobtemplate: q.Jobtemplate.replaceDB(db), - Kaniko: q.Kaniko.replaceDB(db), - ModelDatasetDiscovery: q.ModelDatasetDiscovery.replaceDB(db), - ModelDatasetSource: q.ModelDatasetSource.replaceDB(db), - ModelDownload: q.ModelDownload.replaceDB(db), - ModelDownloadSubmission: q.ModelDownloadSubmission.replaceDB(db), - PrequeueConfig: q.PrequeueConfig.replaceDB(db), - QueueQuotaLimit: q.QueueQuotaLimit.replaceDB(db), - Resource: q.Resource.replaceDB(db), - ResourceNetwork: q.ResourceNetwork.replaceDB(db), - ResourceVGPU: q.ResourceVGPU.replaceDB(db), - SystemConfig: q.SystemConfig.replaceDB(db), - User: q.User.replaceDB(db), - UserAccount: q.UserAccount.replaceDB(db), - UserDataset: q.UserDataset.replaceDB(db), - UserModelDownload: q.UserModelDownload.replaceDB(db), + db: db, + AITask: q.AITask.replaceDB(db), + Account: q.Account.replaceDB(db), + AccountDataset: q.AccountDataset.replaceDB(db), + Alert: q.Alert.replaceDB(db), + ApprovalOrder: q.ApprovalOrder.replaceDB(db), + CronJobConfig: q.CronJobConfig.replaceDB(db), + CronJobRecord: q.CronJobRecord.replaceDB(db), + CudaBaseImage: q.CudaBaseImage.replaceDB(db), + Dataset: q.Dataset.replaceDB(db), + GpuAnalysis: q.GpuAnalysis.replaceDB(db), + Image: q.Image.replaceDB(db), + ImageAccount: q.ImageAccount.replaceDB(db), + ImageUser: q.ImageUser.replaceDB(db), + Job: q.Job.replaceDB(db), + Jobtemplate: q.Jobtemplate.replaceDB(db), + Kaniko: q.Kaniko.replaceDB(db), + ModelDownload: q.ModelDownload.replaceDB(db), + PrequeueConfig: q.PrequeueConfig.replaceDB(db), + QueueQuotaLimit: q.QueueQuotaLimit.replaceDB(db), + Resource: q.Resource.replaceDB(db), + ResourceNetwork: q.ResourceNetwork.replaceDB(db), + ResourceVGPU: q.ResourceVGPU.replaceDB(db), + SystemConfig: q.SystemConfig.replaceDB(db), + User: q.User.replaceDB(db), + UserAccount: q.UserAccount.replaceDB(db), + UserDataset: q.UserDataset.replaceDB(db), + UserModelDownload: q.UserModelDownload.replaceDB(db), } } type queryCtx struct { - AITask IAITaskDo - Account IAccountDo - AccountDataset IAccountDatasetDo - Alert IAlertDo - ApprovalOrder IApprovalOrderDo - CronJobConfig ICronJobConfigDo - CronJobRecord ICronJobRecordDo - CudaBaseImage ICudaBaseImageDo - Dataset IDatasetDo - GpuAnalysis IGpuAnalysisDo - Image IImageDo - ImageAccount IImageAccountDo - ImageUser IImageUserDo - Job IJobDo - Jobtemplate IJobtemplateDo - Kaniko IKanikoDo - ModelDatasetDiscovery IModelDatasetDiscoveryDo - ModelDatasetSource IModelDatasetSourceDo - ModelDownload IModelDownloadDo - ModelDownloadSubmission IModelDownloadSubmissionDo - PrequeueConfig IPrequeueConfigDo - QueueQuotaLimit IQueueQuotaLimitDo - Resource IResourceDo - ResourceNetwork IResourceNetworkDo - ResourceVGPU IResourceVGPUDo - SystemConfig ISystemConfigDo - User IUserDo - UserAccount IUserAccountDo - UserDataset IUserDatasetDo - UserModelDownload IUserModelDownloadDo + AITask IAITaskDo + Account IAccountDo + AccountDataset IAccountDatasetDo + Alert IAlertDo + ApprovalOrder IApprovalOrderDo + CronJobConfig ICronJobConfigDo + CronJobRecord ICronJobRecordDo + CudaBaseImage ICudaBaseImageDo + Dataset IDatasetDo + GpuAnalysis IGpuAnalysisDo + Image IImageDo + ImageAccount IImageAccountDo + ImageUser IImageUserDo + Job IJobDo + Jobtemplate IJobtemplateDo + Kaniko IKanikoDo + ModelDownload IModelDownloadDo + PrequeueConfig IPrequeueConfigDo + QueueQuotaLimit IQueueQuotaLimitDo + Resource IResourceDo + ResourceNetwork IResourceNetworkDo + ResourceVGPU IResourceVGPUDo + SystemConfig ISystemConfigDo + User IUserDo + UserAccount IUserAccountDo + UserDataset IUserDatasetDo + UserModelDownload IUserModelDownloadDo } func (q *Query) WithContext(ctx context.Context) *queryCtx { return &queryCtx{ - AITask: q.AITask.WithContext(ctx), - Account: q.Account.WithContext(ctx), - AccountDataset: q.AccountDataset.WithContext(ctx), - Alert: q.Alert.WithContext(ctx), - ApprovalOrder: q.ApprovalOrder.WithContext(ctx), - CronJobConfig: q.CronJobConfig.WithContext(ctx), - CronJobRecord: q.CronJobRecord.WithContext(ctx), - CudaBaseImage: q.CudaBaseImage.WithContext(ctx), - Dataset: q.Dataset.WithContext(ctx), - GpuAnalysis: q.GpuAnalysis.WithContext(ctx), - Image: q.Image.WithContext(ctx), - ImageAccount: q.ImageAccount.WithContext(ctx), - ImageUser: q.ImageUser.WithContext(ctx), - Job: q.Job.WithContext(ctx), - Jobtemplate: q.Jobtemplate.WithContext(ctx), - Kaniko: q.Kaniko.WithContext(ctx), - ModelDatasetDiscovery: q.ModelDatasetDiscovery.WithContext(ctx), - ModelDatasetSource: q.ModelDatasetSource.WithContext(ctx), - ModelDownload: q.ModelDownload.WithContext(ctx), - ModelDownloadSubmission: q.ModelDownloadSubmission.WithContext(ctx), - PrequeueConfig: q.PrequeueConfig.WithContext(ctx), - QueueQuotaLimit: q.QueueQuotaLimit.WithContext(ctx), - Resource: q.Resource.WithContext(ctx), - ResourceNetwork: q.ResourceNetwork.WithContext(ctx), - ResourceVGPU: q.ResourceVGPU.WithContext(ctx), - SystemConfig: q.SystemConfig.WithContext(ctx), - User: q.User.WithContext(ctx), - UserAccount: q.UserAccount.WithContext(ctx), - UserDataset: q.UserDataset.WithContext(ctx), - UserModelDownload: q.UserModelDownload.WithContext(ctx), + AITask: q.AITask.WithContext(ctx), + Account: q.Account.WithContext(ctx), + AccountDataset: q.AccountDataset.WithContext(ctx), + Alert: q.Alert.WithContext(ctx), + ApprovalOrder: q.ApprovalOrder.WithContext(ctx), + CronJobConfig: q.CronJobConfig.WithContext(ctx), + CronJobRecord: q.CronJobRecord.WithContext(ctx), + CudaBaseImage: q.CudaBaseImage.WithContext(ctx), + Dataset: q.Dataset.WithContext(ctx), + GpuAnalysis: q.GpuAnalysis.WithContext(ctx), + Image: q.Image.WithContext(ctx), + ImageAccount: q.ImageAccount.WithContext(ctx), + ImageUser: q.ImageUser.WithContext(ctx), + Job: q.Job.WithContext(ctx), + Jobtemplate: q.Jobtemplate.WithContext(ctx), + Kaniko: q.Kaniko.WithContext(ctx), + ModelDownload: q.ModelDownload.WithContext(ctx), + PrequeueConfig: q.PrequeueConfig.WithContext(ctx), + QueueQuotaLimit: q.QueueQuotaLimit.WithContext(ctx), + Resource: q.Resource.WithContext(ctx), + ResourceNetwork: q.ResourceNetwork.WithContext(ctx), + ResourceVGPU: q.ResourceVGPU.WithContext(ctx), + SystemConfig: q.SystemConfig.WithContext(ctx), + User: q.User.WithContext(ctx), + UserAccount: q.UserAccount.WithContext(ctx), + UserDataset: q.UserDataset.WithContext(ctx), + UserModelDownload: q.UserModelDownload.WithContext(ctx), } } diff --git a/backend/dao/query/jobs.gen.go b/backend/dao/query/jobs.gen.go index baf6050db..0adf9808e 100644 --- a/backend/dao/query/jobs.gen.go +++ b/backend/dao/query/jobs.gen.go @@ -124,7 +124,7 @@ type job struct { KeepWhenLowResourceUsage field.Bool // 当资源利用率低时是否保留 LockedTimestamp field.Time // 作业锁定时间 LastSettledAt field.Time // 作业上次结算时间 - BilledPointsTotal field.Int64 // 作业累计已结算点数(内部微点) + BilledPointsTotal field.Int64 // 作业累计已结算点数 ProfileData field.Field // 作业的性能数据 ScheduleData field.Field // 作业的调度数据 Events field.Field // 作业的事件 (运行时、失败时采集) @@ -202,7 +202,7 @@ func (j *job) GetFieldByName(fieldName string) (field.OrderExpr, bool) { } func (j *job) fillFieldMap() { - j.fieldMap = make(map[string]field.Expr, 32) + j.fieldMap = make(map[string]field.Expr, 31) j.fieldMap["id"] = j.ID j.fieldMap["created_at"] = j.CreatedAt j.fieldMap["updated_at"] = j.UpdatedAt diff --git a/backend/dao/query/model_dataset_discoveries.gen.go b/backend/dao/query/model_dataset_discoveries.gen.go deleted file mode 100644 index 816d2a8bc..000000000 --- a/backend/dao/query/model_dataset_discoveries.gen.go +++ /dev/null @@ -1,459 +0,0 @@ -// Code generated by gorm.io/gen. DO NOT EDIT. -// Code generated by gorm.io/gen. DO NOT EDIT. -// Code generated by gorm.io/gen. DO NOT EDIT. - -package query - -import ( - "context" - "database/sql" - - "gorm.io/gorm" - "gorm.io/gorm/clause" - "gorm.io/gorm/schema" - - "gorm.io/gen" - "gorm.io/gen/field" - - "gorm.io/plugin/dbresolver" - - "github.com/raids-lab/crater/dao/model" -) - -func newModelDatasetDiscovery(db *gorm.DB, opts ...gen.DOOption) modelDatasetDiscovery { - _modelDatasetDiscovery := modelDatasetDiscovery{} - - _modelDatasetDiscovery.modelDatasetDiscoveryDo.UseDB(db, opts...) - _modelDatasetDiscovery.modelDatasetDiscoveryDo.UseModel(&model.ModelDatasetDiscovery{}) - - tableName := _modelDatasetDiscovery.modelDatasetDiscoveryDo.TableName() - _modelDatasetDiscovery.ALL = field.NewAsterisk(tableName) - _modelDatasetDiscovery.ID = field.NewUint(tableName, "id") - _modelDatasetDiscovery.CreatedAt = field.NewTime(tableName, "created_at") - _modelDatasetDiscovery.UpdatedAt = field.NewTime(tableName, "updated_at") - _modelDatasetDiscovery.DeletedAt = field.NewField(tableName, "deleted_at") - _modelDatasetDiscovery.DiscoveryKey = field.NewString(tableName, "discovery_key") - _modelDatasetDiscovery.Path = field.NewString(tableName, "path") - _modelDatasetDiscovery.Scope = field.NewString(tableName, "scope") - _modelDatasetDiscovery.ScopeID = field.NewUint(tableName, "scope_id") - _modelDatasetDiscovery.DetectedType = field.NewString(tableName, "detected_type") - _modelDatasetDiscovery.DetectedName = field.NewString(tableName, "detected_name") - _modelDatasetDiscovery.Evidence = field.NewField(tableName, "evidence") - _modelDatasetDiscovery.SizeBytes = field.NewInt64(tableName, "size_bytes") - _modelDatasetDiscovery.DatasetID = field.NewUint(tableName, "dataset_id") - _modelDatasetDiscovery.SourceID = field.NewUint(tableName, "source_id") - _modelDatasetDiscovery.Status = field.NewString(tableName, "status") - _modelDatasetDiscovery.FirstSeenAt = field.NewTime(tableName, "first_seen_at") - _modelDatasetDiscovery.LastSeenAt = field.NewTime(tableName, "last_seen_at") - - _modelDatasetDiscovery.fillFieldMap() - - return _modelDatasetDiscovery -} - -type modelDatasetDiscovery struct { - modelDatasetDiscoveryDo modelDatasetDiscoveryDo - - ALL field.Asterisk - ID field.Uint - CreatedAt field.Time - UpdatedAt field.Time - DeletedAt field.Field - DiscoveryKey field.String // 稳定发现键 - Path field.String // 文件系统路径,允许为空 - Scope field.String - ScopeID field.Uint // 用户或队列ID - DetectedType field.String - DetectedName field.String - Evidence field.Field // 文件系统检测依据 - SizeBytes field.Int64 - DatasetID field.Uint - SourceID field.Uint - Status field.String - FirstSeenAt field.Time - LastSeenAt field.Time - - fieldMap map[string]field.Expr -} - -func (m modelDatasetDiscovery) Table(newTableName string) *modelDatasetDiscovery { - m.modelDatasetDiscoveryDo.UseTable(newTableName) - return m.updateTableName(newTableName) -} - -func (m modelDatasetDiscovery) As(alias string) *modelDatasetDiscovery { - m.modelDatasetDiscoveryDo.DO = *(m.modelDatasetDiscoveryDo.As(alias).(*gen.DO)) - return m.updateTableName(alias) -} - -func (m *modelDatasetDiscovery) updateTableName(table string) *modelDatasetDiscovery { - m.ALL = field.NewAsterisk(table) - m.ID = field.NewUint(table, "id") - m.CreatedAt = field.NewTime(table, "created_at") - m.UpdatedAt = field.NewTime(table, "updated_at") - m.DeletedAt = field.NewField(table, "deleted_at") - m.DiscoveryKey = field.NewString(table, "discovery_key") - m.Path = field.NewString(table, "path") - m.Scope = field.NewString(table, "scope") - m.ScopeID = field.NewUint(table, "scope_id") - m.DetectedType = field.NewString(table, "detected_type") - m.DetectedName = field.NewString(table, "detected_name") - m.Evidence = field.NewField(table, "evidence") - m.SizeBytes = field.NewInt64(table, "size_bytes") - m.DatasetID = field.NewUint(table, "dataset_id") - m.SourceID = field.NewUint(table, "source_id") - m.Status = field.NewString(table, "status") - m.FirstSeenAt = field.NewTime(table, "first_seen_at") - m.LastSeenAt = field.NewTime(table, "last_seen_at") - - m.fillFieldMap() - - return m -} - -func (m *modelDatasetDiscovery) WithContext(ctx context.Context) IModelDatasetDiscoveryDo { - return m.modelDatasetDiscoveryDo.WithContext(ctx) -} - -func (m modelDatasetDiscovery) TableName() string { return m.modelDatasetDiscoveryDo.TableName() } - -func (m modelDatasetDiscovery) Alias() string { return m.modelDatasetDiscoveryDo.Alias() } - -func (m modelDatasetDiscovery) Columns(cols ...field.Expr) gen.Columns { - return m.modelDatasetDiscoveryDo.Columns(cols...) -} - -func (m *modelDatasetDiscovery) GetFieldByName(fieldName string) (field.OrderExpr, bool) { - _f, ok := m.fieldMap[fieldName] - if !ok || _f == nil { - return nil, false - } - _oe, ok := _f.(field.OrderExpr) - return _oe, ok -} - -func (m *modelDatasetDiscovery) fillFieldMap() { - m.fieldMap = make(map[string]field.Expr, 17) - m.fieldMap["id"] = m.ID - m.fieldMap["created_at"] = m.CreatedAt - m.fieldMap["updated_at"] = m.UpdatedAt - m.fieldMap["deleted_at"] = m.DeletedAt - m.fieldMap["discovery_key"] = m.DiscoveryKey - m.fieldMap["path"] = m.Path - m.fieldMap["scope"] = m.Scope - m.fieldMap["scope_id"] = m.ScopeID - m.fieldMap["detected_type"] = m.DetectedType - m.fieldMap["detected_name"] = m.DetectedName - m.fieldMap["evidence"] = m.Evidence - m.fieldMap["size_bytes"] = m.SizeBytes - m.fieldMap["dataset_id"] = m.DatasetID - m.fieldMap["source_id"] = m.SourceID - m.fieldMap["status"] = m.Status - m.fieldMap["first_seen_at"] = m.FirstSeenAt - m.fieldMap["last_seen_at"] = m.LastSeenAt -} - -func (m modelDatasetDiscovery) clone(db *gorm.DB) modelDatasetDiscovery { - m.modelDatasetDiscoveryDo.ReplaceConnPool(db.Statement.ConnPool) - return m -} - -func (m modelDatasetDiscovery) replaceDB(db *gorm.DB) modelDatasetDiscovery { - m.modelDatasetDiscoveryDo.ReplaceDB(db) - return m -} - -type modelDatasetDiscoveryDo struct{ gen.DO } - -type IModelDatasetDiscoveryDo interface { - gen.SubQuery - Debug() IModelDatasetDiscoveryDo - WithContext(ctx context.Context) IModelDatasetDiscoveryDo - WithResult(fc func(tx gen.Dao)) gen.ResultInfo - ReplaceDB(db *gorm.DB) - ReadDB() IModelDatasetDiscoveryDo - WriteDB() IModelDatasetDiscoveryDo - As(alias string) gen.Dao - Session(config *gorm.Session) IModelDatasetDiscoveryDo - Columns(cols ...field.Expr) gen.Columns - Clauses(conds ...clause.Expression) IModelDatasetDiscoveryDo - Not(conds ...gen.Condition) IModelDatasetDiscoveryDo - Or(conds ...gen.Condition) IModelDatasetDiscoveryDo - Select(conds ...field.Expr) IModelDatasetDiscoveryDo - Where(conds ...gen.Condition) IModelDatasetDiscoveryDo - Order(conds ...field.Expr) IModelDatasetDiscoveryDo - Distinct(cols ...field.Expr) IModelDatasetDiscoveryDo - Omit(cols ...field.Expr) IModelDatasetDiscoveryDo - Join(table schema.Tabler, on ...field.Expr) IModelDatasetDiscoveryDo - LeftJoin(table schema.Tabler, on ...field.Expr) IModelDatasetDiscoveryDo - RightJoin(table schema.Tabler, on ...field.Expr) IModelDatasetDiscoveryDo - Group(cols ...field.Expr) IModelDatasetDiscoveryDo - Having(conds ...gen.Condition) IModelDatasetDiscoveryDo - Limit(limit int) IModelDatasetDiscoveryDo - Offset(offset int) IModelDatasetDiscoveryDo - Count() (count int64, err error) - Scopes(funcs ...func(gen.Dao) gen.Dao) IModelDatasetDiscoveryDo - Unscoped() IModelDatasetDiscoveryDo - Create(values ...*model.ModelDatasetDiscovery) error - CreateInBatches(values []*model.ModelDatasetDiscovery, batchSize int) error - Save(values ...*model.ModelDatasetDiscovery) error - First() (*model.ModelDatasetDiscovery, error) - Take() (*model.ModelDatasetDiscovery, error) - Last() (*model.ModelDatasetDiscovery, error) - Find() ([]*model.ModelDatasetDiscovery, error) - FindInBatch(batchSize int, fc func(tx gen.Dao, batch int) error) (results []*model.ModelDatasetDiscovery, err error) - FindInBatches(result *[]*model.ModelDatasetDiscovery, batchSize int, fc func(tx gen.Dao, batch int) error) error - Pluck(column field.Expr, dest interface{}) error - Delete(...*model.ModelDatasetDiscovery) (info gen.ResultInfo, err error) - Update(column field.Expr, value interface{}) (info gen.ResultInfo, err error) - UpdateSimple(columns ...field.AssignExpr) (info gen.ResultInfo, err error) - Updates(value interface{}) (info gen.ResultInfo, err error) - UpdateColumn(column field.Expr, value interface{}) (info gen.ResultInfo, err error) - UpdateColumnSimple(columns ...field.AssignExpr) (info gen.ResultInfo, err error) - UpdateColumns(value interface{}) (info gen.ResultInfo, err error) - UpdateFrom(q gen.SubQuery) gen.Dao - Attrs(attrs ...field.AssignExpr) IModelDatasetDiscoveryDo - Assign(attrs ...field.AssignExpr) IModelDatasetDiscoveryDo - Joins(fields ...field.RelationField) IModelDatasetDiscoveryDo - Preload(fields ...field.RelationField) IModelDatasetDiscoveryDo - FirstOrInit() (*model.ModelDatasetDiscovery, error) - FirstOrCreate() (*model.ModelDatasetDiscovery, error) - FindByPage(offset int, limit int) (result []*model.ModelDatasetDiscovery, count int64, err error) - ScanByPage(result interface{}, offset int, limit int) (count int64, err error) - Rows() (*sql.Rows, error) - Row() *sql.Row - Scan(result interface{}) (err error) - Returning(value interface{}, columns ...string) IModelDatasetDiscoveryDo - UnderlyingDB() *gorm.DB - schema.Tabler -} - -func (m modelDatasetDiscoveryDo) Debug() IModelDatasetDiscoveryDo { - return m.withDO(m.DO.Debug()) -} - -func (m modelDatasetDiscoveryDo) WithContext(ctx context.Context) IModelDatasetDiscoveryDo { - return m.withDO(m.DO.WithContext(ctx)) -} - -func (m modelDatasetDiscoveryDo) ReadDB() IModelDatasetDiscoveryDo { - return m.Clauses(dbresolver.Read) -} - -func (m modelDatasetDiscoveryDo) WriteDB() IModelDatasetDiscoveryDo { - return m.Clauses(dbresolver.Write) -} - -func (m modelDatasetDiscoveryDo) Session(config *gorm.Session) IModelDatasetDiscoveryDo { - return m.withDO(m.DO.Session(config)) -} - -func (m modelDatasetDiscoveryDo) Clauses(conds ...clause.Expression) IModelDatasetDiscoveryDo { - return m.withDO(m.DO.Clauses(conds...)) -} - -func (m modelDatasetDiscoveryDo) Returning(value interface{}, columns ...string) IModelDatasetDiscoveryDo { - return m.withDO(m.DO.Returning(value, columns...)) -} - -func (m modelDatasetDiscoveryDo) Not(conds ...gen.Condition) IModelDatasetDiscoveryDo { - return m.withDO(m.DO.Not(conds...)) -} - -func (m modelDatasetDiscoveryDo) Or(conds ...gen.Condition) IModelDatasetDiscoveryDo { - return m.withDO(m.DO.Or(conds...)) -} - -func (m modelDatasetDiscoveryDo) Select(conds ...field.Expr) IModelDatasetDiscoveryDo { - return m.withDO(m.DO.Select(conds...)) -} - -func (m modelDatasetDiscoveryDo) Where(conds ...gen.Condition) IModelDatasetDiscoveryDo { - return m.withDO(m.DO.Where(conds...)) -} - -func (m modelDatasetDiscoveryDo) Order(conds ...field.Expr) IModelDatasetDiscoveryDo { - return m.withDO(m.DO.Order(conds...)) -} - -func (m modelDatasetDiscoveryDo) Distinct(cols ...field.Expr) IModelDatasetDiscoveryDo { - return m.withDO(m.DO.Distinct(cols...)) -} - -func (m modelDatasetDiscoveryDo) Omit(cols ...field.Expr) IModelDatasetDiscoveryDo { - return m.withDO(m.DO.Omit(cols...)) -} - -func (m modelDatasetDiscoveryDo) Join(table schema.Tabler, on ...field.Expr) IModelDatasetDiscoveryDo { - return m.withDO(m.DO.Join(table, on...)) -} - -func (m modelDatasetDiscoveryDo) LeftJoin(table schema.Tabler, on ...field.Expr) IModelDatasetDiscoveryDo { - return m.withDO(m.DO.LeftJoin(table, on...)) -} - -func (m modelDatasetDiscoveryDo) RightJoin(table schema.Tabler, on ...field.Expr) IModelDatasetDiscoveryDo { - return m.withDO(m.DO.RightJoin(table, on...)) -} - -func (m modelDatasetDiscoveryDo) Group(cols ...field.Expr) IModelDatasetDiscoveryDo { - return m.withDO(m.DO.Group(cols...)) -} - -func (m modelDatasetDiscoveryDo) Having(conds ...gen.Condition) IModelDatasetDiscoveryDo { - return m.withDO(m.DO.Having(conds...)) -} - -func (m modelDatasetDiscoveryDo) Limit(limit int) IModelDatasetDiscoveryDo { - return m.withDO(m.DO.Limit(limit)) -} - -func (m modelDatasetDiscoveryDo) Offset(offset int) IModelDatasetDiscoveryDo { - return m.withDO(m.DO.Offset(offset)) -} - -func (m modelDatasetDiscoveryDo) Scopes(funcs ...func(gen.Dao) gen.Dao) IModelDatasetDiscoveryDo { - return m.withDO(m.DO.Scopes(funcs...)) -} - -func (m modelDatasetDiscoveryDo) Unscoped() IModelDatasetDiscoveryDo { - return m.withDO(m.DO.Unscoped()) -} - -func (m modelDatasetDiscoveryDo) Create(values ...*model.ModelDatasetDiscovery) error { - if len(values) == 0 { - return nil - } - return m.DO.Create(values) -} - -func (m modelDatasetDiscoveryDo) CreateInBatches(values []*model.ModelDatasetDiscovery, batchSize int) error { - return m.DO.CreateInBatches(values, batchSize) -} - -// Save : !!! underlying implementation is different with GORM -// The method is equivalent to executing the statement: db.Clauses(clause.OnConflict{UpdateAll: true}).Create(values) -func (m modelDatasetDiscoveryDo) Save(values ...*model.ModelDatasetDiscovery) error { - if len(values) == 0 { - return nil - } - return m.DO.Save(values) -} - -func (m modelDatasetDiscoveryDo) First() (*model.ModelDatasetDiscovery, error) { - if result, err := m.DO.First(); err != nil { - return nil, err - } else { - return result.(*model.ModelDatasetDiscovery), nil - } -} - -func (m modelDatasetDiscoveryDo) Take() (*model.ModelDatasetDiscovery, error) { - if result, err := m.DO.Take(); err != nil { - return nil, err - } else { - return result.(*model.ModelDatasetDiscovery), nil - } -} - -func (m modelDatasetDiscoveryDo) Last() (*model.ModelDatasetDiscovery, error) { - if result, err := m.DO.Last(); err != nil { - return nil, err - } else { - return result.(*model.ModelDatasetDiscovery), nil - } -} - -func (m modelDatasetDiscoveryDo) Find() ([]*model.ModelDatasetDiscovery, error) { - result, err := m.DO.Find() - return result.([]*model.ModelDatasetDiscovery), err -} - -func (m modelDatasetDiscoveryDo) FindInBatch(batchSize int, fc func(tx gen.Dao, batch int) error) (results []*model.ModelDatasetDiscovery, err error) { - buf := make([]*model.ModelDatasetDiscovery, 0, batchSize) - err = m.DO.FindInBatches(&buf, batchSize, func(tx gen.Dao, batch int) error { - defer func() { results = append(results, buf...) }() - return fc(tx, batch) - }) - return results, err -} - -func (m modelDatasetDiscoveryDo) FindInBatches(result *[]*model.ModelDatasetDiscovery, batchSize int, fc func(tx gen.Dao, batch int) error) error { - return m.DO.FindInBatches(result, batchSize, fc) -} - -func (m modelDatasetDiscoveryDo) Attrs(attrs ...field.AssignExpr) IModelDatasetDiscoveryDo { - return m.withDO(m.DO.Attrs(attrs...)) -} - -func (m modelDatasetDiscoveryDo) Assign(attrs ...field.AssignExpr) IModelDatasetDiscoveryDo { - return m.withDO(m.DO.Assign(attrs...)) -} - -func (m modelDatasetDiscoveryDo) Joins(fields ...field.RelationField) IModelDatasetDiscoveryDo { - for _, _f := range fields { - m = *m.withDO(m.DO.Joins(_f)) - } - return &m -} - -func (m modelDatasetDiscoveryDo) Preload(fields ...field.RelationField) IModelDatasetDiscoveryDo { - for _, _f := range fields { - m = *m.withDO(m.DO.Preload(_f)) - } - return &m -} - -func (m modelDatasetDiscoveryDo) FirstOrInit() (*model.ModelDatasetDiscovery, error) { - if result, err := m.DO.FirstOrInit(); err != nil { - return nil, err - } else { - return result.(*model.ModelDatasetDiscovery), nil - } -} - -func (m modelDatasetDiscoveryDo) FirstOrCreate() (*model.ModelDatasetDiscovery, error) { - if result, err := m.DO.FirstOrCreate(); err != nil { - return nil, err - } else { - return result.(*model.ModelDatasetDiscovery), nil - } -} - -func (m modelDatasetDiscoveryDo) FindByPage(offset int, limit int) (result []*model.ModelDatasetDiscovery, count int64, err error) { - result, err = m.Offset(offset).Limit(limit).Find() - if err != nil { - return - } - - if size := len(result); 0 < limit && 0 < size && size < limit { - count = int64(size + offset) - return - } - - count, err = m.Offset(-1).Limit(-1).Count() - return -} - -func (m modelDatasetDiscoveryDo) ScanByPage(result interface{}, offset int, limit int) (count int64, err error) { - count, err = m.Count() - if err != nil { - return - } - - err = m.Offset(offset).Limit(limit).Scan(result) - return -} - -func (m modelDatasetDiscoveryDo) Scan(result interface{}) (err error) { - return m.DO.Scan(result) -} - -func (m modelDatasetDiscoveryDo) Delete(models ...*model.ModelDatasetDiscovery) (result gen.ResultInfo, err error) { - return m.DO.Delete(models) -} - -func (m *modelDatasetDiscoveryDo) withDO(do gen.Dao) *modelDatasetDiscoveryDo { - m.DO = *do.(*gen.DO) - return m -} diff --git a/backend/dao/query/model_dataset_sources.gen.go b/backend/dao/query/model_dataset_sources.gen.go deleted file mode 100644 index 988e75e20..000000000 --- a/backend/dao/query/model_dataset_sources.gen.go +++ /dev/null @@ -1,503 +0,0 @@ -// Code generated by gorm.io/gen. DO NOT EDIT. -// Code generated by gorm.io/gen. DO NOT EDIT. -// Code generated by gorm.io/gen. DO NOT EDIT. - -package query - -import ( - "context" - "database/sql" - - "gorm.io/gorm" - "gorm.io/gorm/clause" - "gorm.io/gorm/schema" - - "gorm.io/gen" - "gorm.io/gen/field" - - "gorm.io/plugin/dbresolver" - - "github.com/raids-lab/crater/dao/model" -) - -func newModelDatasetSource(db *gorm.DB, opts ...gen.DOOption) modelDatasetSource { - _modelDatasetSource := modelDatasetSource{} - - _modelDatasetSource.modelDatasetSourceDo.UseDB(db, opts...) - _modelDatasetSource.modelDatasetSourceDo.UseModel(&model.ModelDatasetSource{}) - - tableName := _modelDatasetSource.modelDatasetSourceDo.TableName() - _modelDatasetSource.ALL = field.NewAsterisk(tableName) - _modelDatasetSource.ID = field.NewUint(tableName, "id") - _modelDatasetSource.CreatedAt = field.NewTime(tableName, "created_at") - _modelDatasetSource.UpdatedAt = field.NewTime(tableName, "updated_at") - _modelDatasetSource.DeletedAt = field.NewField(tableName, "deleted_at") - _modelDatasetSource.Provider = field.NewString(tableName, "provider") - _modelDatasetSource.ResourceType = field.NewString(tableName, "resource_type") - _modelDatasetSource.RepositoryID = field.NewString(tableName, "repository_id") - _modelDatasetSource.RepositoryURL = field.NewString(tableName, "repository_url") - _modelDatasetSource.Organization = field.NewString(tableName, "organization") - _modelDatasetSource.LogoURL = field.NewString(tableName, "logo_url") - _modelDatasetSource.LogoData = field.NewBytes(tableName, "logo_data") - _modelDatasetSource.LogoContentType = field.NewString(tableName, "logo_content_type") - _modelDatasetSource.DisplayName = field.NewString(tableName, "display_name") - _modelDatasetSource.Description = field.NewString(tableName, "description") - _modelDatasetSource.Readme = field.NewString(tableName, "readme") - _modelDatasetSource.License = field.NewString(tableName, "license") - _modelDatasetSource.Task = field.NewString(tableName, "task") - _modelDatasetSource.Library = field.NewString(tableName, "library") - _modelDatasetSource.ModelType = field.NewString(tableName, "model_type") - _modelDatasetSource.ParameterCount = field.NewInt64(tableName, "parameter_count") - _modelDatasetSource.Private = field.NewBool(tableName, "private") - _modelDatasetSource.Gated = field.NewBool(tableName, "gated") - _modelDatasetSource.LoginRequired = field.NewBool(tableName, "login_required") - _modelDatasetSource.Downloads = field.NewInt64(tableName, "downloads") - _modelDatasetSource.Likes = field.NewInt64(tableName, "likes") - _modelDatasetSource.SourceCreatedAt = field.NewTime(tableName, "source_created_at") - _modelDatasetSource.SourceUpdatedAt = field.NewTime(tableName, "source_updated_at") - _modelDatasetSource.MetadataRefreshedAt = field.NewTime(tableName, "metadata_refreshed_at") - - _modelDatasetSource.fillFieldMap() - - return _modelDatasetSource -} - -type modelDatasetSource struct { - modelDatasetSourceDo modelDatasetSourceDo - - ALL field.Asterisk - ID field.Uint - CreatedAt field.Time - UpdatedAt field.Time - DeletedAt field.Field - Provider field.String - ResourceType field.String - RepositoryID field.String - RepositoryURL field.String // 外部仓库页面地址 - Organization field.String // 源站组织或作者 - LogoURL field.String // 源站组织头像地址 - LogoData field.Bytes // 平台缓存的源站头像 - LogoContentType field.String // 平台缓存头像的Content-Type - DisplayName field.String // 源站展示名称 - Description field.String // 源站简介摘要 - Readme field.String // 源站README内容(截断保存) - License field.String // 源站许可证 - Task field.String // 源站任务分类 - Library field.String // 源站框架或库 - ModelType field.String // 源站模型类型 - ParameterCount field.Int64 // 模型参数量 - Private field.Bool // 源站是否私有 - Gated field.Bool // 源站是否需要申请访问 - LoginRequired field.Bool // 源站是否要求登录下载 - Downloads field.Int64 // 源站下载次数 - Likes field.Int64 // 源站点赞次数 - SourceCreatedAt field.Time // 源站创建时间 - SourceUpdatedAt field.Time // 源站更新时间 - MetadataRefreshedAt field.Time // 源站元数据刷新时间 - - fieldMap map[string]field.Expr -} - -func (m modelDatasetSource) Table(newTableName string) *modelDatasetSource { - m.modelDatasetSourceDo.UseTable(newTableName) - return m.updateTableName(newTableName) -} - -func (m modelDatasetSource) As(alias string) *modelDatasetSource { - m.modelDatasetSourceDo.DO = *(m.modelDatasetSourceDo.As(alias).(*gen.DO)) - return m.updateTableName(alias) -} - -func (m *modelDatasetSource) updateTableName(table string) *modelDatasetSource { - m.ALL = field.NewAsterisk(table) - m.ID = field.NewUint(table, "id") - m.CreatedAt = field.NewTime(table, "created_at") - m.UpdatedAt = field.NewTime(table, "updated_at") - m.DeletedAt = field.NewField(table, "deleted_at") - m.Provider = field.NewString(table, "provider") - m.ResourceType = field.NewString(table, "resource_type") - m.RepositoryID = field.NewString(table, "repository_id") - m.RepositoryURL = field.NewString(table, "repository_url") - m.Organization = field.NewString(table, "organization") - m.LogoURL = field.NewString(table, "logo_url") - m.LogoData = field.NewBytes(table, "logo_data") - m.LogoContentType = field.NewString(table, "logo_content_type") - m.DisplayName = field.NewString(table, "display_name") - m.Description = field.NewString(table, "description") - m.Readme = field.NewString(table, "readme") - m.License = field.NewString(table, "license") - m.Task = field.NewString(table, "task") - m.Library = field.NewString(table, "library") - m.ModelType = field.NewString(table, "model_type") - m.ParameterCount = field.NewInt64(table, "parameter_count") - m.Private = field.NewBool(table, "private") - m.Gated = field.NewBool(table, "gated") - m.LoginRequired = field.NewBool(table, "login_required") - m.Downloads = field.NewInt64(table, "downloads") - m.Likes = field.NewInt64(table, "likes") - m.SourceCreatedAt = field.NewTime(table, "source_created_at") - m.SourceUpdatedAt = field.NewTime(table, "source_updated_at") - m.MetadataRefreshedAt = field.NewTime(table, "metadata_refreshed_at") - - m.fillFieldMap() - - return m -} - -func (m *modelDatasetSource) WithContext(ctx context.Context) IModelDatasetSourceDo { - return m.modelDatasetSourceDo.WithContext(ctx) -} - -func (m modelDatasetSource) TableName() string { return m.modelDatasetSourceDo.TableName() } - -func (m modelDatasetSource) Alias() string { return m.modelDatasetSourceDo.Alias() } - -func (m modelDatasetSource) Columns(cols ...field.Expr) gen.Columns { - return m.modelDatasetSourceDo.Columns(cols...) -} - -func (m *modelDatasetSource) GetFieldByName(fieldName string) (field.OrderExpr, bool) { - _f, ok := m.fieldMap[fieldName] - if !ok || _f == nil { - return nil, false - } - _oe, ok := _f.(field.OrderExpr) - return _oe, ok -} - -func (m *modelDatasetSource) fillFieldMap() { - m.fieldMap = make(map[string]field.Expr, 28) - m.fieldMap["id"] = m.ID - m.fieldMap["created_at"] = m.CreatedAt - m.fieldMap["updated_at"] = m.UpdatedAt - m.fieldMap["deleted_at"] = m.DeletedAt - m.fieldMap["provider"] = m.Provider - m.fieldMap["resource_type"] = m.ResourceType - m.fieldMap["repository_id"] = m.RepositoryID - m.fieldMap["repository_url"] = m.RepositoryURL - m.fieldMap["organization"] = m.Organization - m.fieldMap["logo_url"] = m.LogoURL - m.fieldMap["logo_data"] = m.LogoData - m.fieldMap["logo_content_type"] = m.LogoContentType - m.fieldMap["display_name"] = m.DisplayName - m.fieldMap["description"] = m.Description - m.fieldMap["readme"] = m.Readme - m.fieldMap["license"] = m.License - m.fieldMap["task"] = m.Task - m.fieldMap["library"] = m.Library - m.fieldMap["model_type"] = m.ModelType - m.fieldMap["parameter_count"] = m.ParameterCount - m.fieldMap["private"] = m.Private - m.fieldMap["gated"] = m.Gated - m.fieldMap["login_required"] = m.LoginRequired - m.fieldMap["downloads"] = m.Downloads - m.fieldMap["likes"] = m.Likes - m.fieldMap["source_created_at"] = m.SourceCreatedAt - m.fieldMap["source_updated_at"] = m.SourceUpdatedAt - m.fieldMap["metadata_refreshed_at"] = m.MetadataRefreshedAt -} - -func (m modelDatasetSource) clone(db *gorm.DB) modelDatasetSource { - m.modelDatasetSourceDo.ReplaceConnPool(db.Statement.ConnPool) - return m -} - -func (m modelDatasetSource) replaceDB(db *gorm.DB) modelDatasetSource { - m.modelDatasetSourceDo.ReplaceDB(db) - return m -} - -type modelDatasetSourceDo struct{ gen.DO } - -type IModelDatasetSourceDo interface { - gen.SubQuery - Debug() IModelDatasetSourceDo - WithContext(ctx context.Context) IModelDatasetSourceDo - WithResult(fc func(tx gen.Dao)) gen.ResultInfo - ReplaceDB(db *gorm.DB) - ReadDB() IModelDatasetSourceDo - WriteDB() IModelDatasetSourceDo - As(alias string) gen.Dao - Session(config *gorm.Session) IModelDatasetSourceDo - Columns(cols ...field.Expr) gen.Columns - Clauses(conds ...clause.Expression) IModelDatasetSourceDo - Not(conds ...gen.Condition) IModelDatasetSourceDo - Or(conds ...gen.Condition) IModelDatasetSourceDo - Select(conds ...field.Expr) IModelDatasetSourceDo - Where(conds ...gen.Condition) IModelDatasetSourceDo - Order(conds ...field.Expr) IModelDatasetSourceDo - Distinct(cols ...field.Expr) IModelDatasetSourceDo - Omit(cols ...field.Expr) IModelDatasetSourceDo - Join(table schema.Tabler, on ...field.Expr) IModelDatasetSourceDo - LeftJoin(table schema.Tabler, on ...field.Expr) IModelDatasetSourceDo - RightJoin(table schema.Tabler, on ...field.Expr) IModelDatasetSourceDo - Group(cols ...field.Expr) IModelDatasetSourceDo - Having(conds ...gen.Condition) IModelDatasetSourceDo - Limit(limit int) IModelDatasetSourceDo - Offset(offset int) IModelDatasetSourceDo - Count() (count int64, err error) - Scopes(funcs ...func(gen.Dao) gen.Dao) IModelDatasetSourceDo - Unscoped() IModelDatasetSourceDo - Create(values ...*model.ModelDatasetSource) error - CreateInBatches(values []*model.ModelDatasetSource, batchSize int) error - Save(values ...*model.ModelDatasetSource) error - First() (*model.ModelDatasetSource, error) - Take() (*model.ModelDatasetSource, error) - Last() (*model.ModelDatasetSource, error) - Find() ([]*model.ModelDatasetSource, error) - FindInBatch(batchSize int, fc func(tx gen.Dao, batch int) error) (results []*model.ModelDatasetSource, err error) - FindInBatches(result *[]*model.ModelDatasetSource, batchSize int, fc func(tx gen.Dao, batch int) error) error - Pluck(column field.Expr, dest interface{}) error - Delete(...*model.ModelDatasetSource) (info gen.ResultInfo, err error) - Update(column field.Expr, value interface{}) (info gen.ResultInfo, err error) - UpdateSimple(columns ...field.AssignExpr) (info gen.ResultInfo, err error) - Updates(value interface{}) (info gen.ResultInfo, err error) - UpdateColumn(column field.Expr, value interface{}) (info gen.ResultInfo, err error) - UpdateColumnSimple(columns ...field.AssignExpr) (info gen.ResultInfo, err error) - UpdateColumns(value interface{}) (info gen.ResultInfo, err error) - UpdateFrom(q gen.SubQuery) gen.Dao - Attrs(attrs ...field.AssignExpr) IModelDatasetSourceDo - Assign(attrs ...field.AssignExpr) IModelDatasetSourceDo - Joins(fields ...field.RelationField) IModelDatasetSourceDo - Preload(fields ...field.RelationField) IModelDatasetSourceDo - FirstOrInit() (*model.ModelDatasetSource, error) - FirstOrCreate() (*model.ModelDatasetSource, error) - FindByPage(offset int, limit int) (result []*model.ModelDatasetSource, count int64, err error) - ScanByPage(result interface{}, offset int, limit int) (count int64, err error) - Rows() (*sql.Rows, error) - Row() *sql.Row - Scan(result interface{}) (err error) - Returning(value interface{}, columns ...string) IModelDatasetSourceDo - UnderlyingDB() *gorm.DB - schema.Tabler -} - -func (m modelDatasetSourceDo) Debug() IModelDatasetSourceDo { - return m.withDO(m.DO.Debug()) -} - -func (m modelDatasetSourceDo) WithContext(ctx context.Context) IModelDatasetSourceDo { - return m.withDO(m.DO.WithContext(ctx)) -} - -func (m modelDatasetSourceDo) ReadDB() IModelDatasetSourceDo { - return m.Clauses(dbresolver.Read) -} - -func (m modelDatasetSourceDo) WriteDB() IModelDatasetSourceDo { - return m.Clauses(dbresolver.Write) -} - -func (m modelDatasetSourceDo) Session(config *gorm.Session) IModelDatasetSourceDo { - return m.withDO(m.DO.Session(config)) -} - -func (m modelDatasetSourceDo) Clauses(conds ...clause.Expression) IModelDatasetSourceDo { - return m.withDO(m.DO.Clauses(conds...)) -} - -func (m modelDatasetSourceDo) Returning(value interface{}, columns ...string) IModelDatasetSourceDo { - return m.withDO(m.DO.Returning(value, columns...)) -} - -func (m modelDatasetSourceDo) Not(conds ...gen.Condition) IModelDatasetSourceDo { - return m.withDO(m.DO.Not(conds...)) -} - -func (m modelDatasetSourceDo) Or(conds ...gen.Condition) IModelDatasetSourceDo { - return m.withDO(m.DO.Or(conds...)) -} - -func (m modelDatasetSourceDo) Select(conds ...field.Expr) IModelDatasetSourceDo { - return m.withDO(m.DO.Select(conds...)) -} - -func (m modelDatasetSourceDo) Where(conds ...gen.Condition) IModelDatasetSourceDo { - return m.withDO(m.DO.Where(conds...)) -} - -func (m modelDatasetSourceDo) Order(conds ...field.Expr) IModelDatasetSourceDo { - return m.withDO(m.DO.Order(conds...)) -} - -func (m modelDatasetSourceDo) Distinct(cols ...field.Expr) IModelDatasetSourceDo { - return m.withDO(m.DO.Distinct(cols...)) -} - -func (m modelDatasetSourceDo) Omit(cols ...field.Expr) IModelDatasetSourceDo { - return m.withDO(m.DO.Omit(cols...)) -} - -func (m modelDatasetSourceDo) Join(table schema.Tabler, on ...field.Expr) IModelDatasetSourceDo { - return m.withDO(m.DO.Join(table, on...)) -} - -func (m modelDatasetSourceDo) LeftJoin(table schema.Tabler, on ...field.Expr) IModelDatasetSourceDo { - return m.withDO(m.DO.LeftJoin(table, on...)) -} - -func (m modelDatasetSourceDo) RightJoin(table schema.Tabler, on ...field.Expr) IModelDatasetSourceDo { - return m.withDO(m.DO.RightJoin(table, on...)) -} - -func (m modelDatasetSourceDo) Group(cols ...field.Expr) IModelDatasetSourceDo { - return m.withDO(m.DO.Group(cols...)) -} - -func (m modelDatasetSourceDo) Having(conds ...gen.Condition) IModelDatasetSourceDo { - return m.withDO(m.DO.Having(conds...)) -} - -func (m modelDatasetSourceDo) Limit(limit int) IModelDatasetSourceDo { - return m.withDO(m.DO.Limit(limit)) -} - -func (m modelDatasetSourceDo) Offset(offset int) IModelDatasetSourceDo { - return m.withDO(m.DO.Offset(offset)) -} - -func (m modelDatasetSourceDo) Scopes(funcs ...func(gen.Dao) gen.Dao) IModelDatasetSourceDo { - return m.withDO(m.DO.Scopes(funcs...)) -} - -func (m modelDatasetSourceDo) Unscoped() IModelDatasetSourceDo { - return m.withDO(m.DO.Unscoped()) -} - -func (m modelDatasetSourceDo) Create(values ...*model.ModelDatasetSource) error { - if len(values) == 0 { - return nil - } - return m.DO.Create(values) -} - -func (m modelDatasetSourceDo) CreateInBatches(values []*model.ModelDatasetSource, batchSize int) error { - return m.DO.CreateInBatches(values, batchSize) -} - -// Save : !!! underlying implementation is different with GORM -// The method is equivalent to executing the statement: db.Clauses(clause.OnConflict{UpdateAll: true}).Create(values) -func (m modelDatasetSourceDo) Save(values ...*model.ModelDatasetSource) error { - if len(values) == 0 { - return nil - } - return m.DO.Save(values) -} - -func (m modelDatasetSourceDo) First() (*model.ModelDatasetSource, error) { - if result, err := m.DO.First(); err != nil { - return nil, err - } else { - return result.(*model.ModelDatasetSource), nil - } -} - -func (m modelDatasetSourceDo) Take() (*model.ModelDatasetSource, error) { - if result, err := m.DO.Take(); err != nil { - return nil, err - } else { - return result.(*model.ModelDatasetSource), nil - } -} - -func (m modelDatasetSourceDo) Last() (*model.ModelDatasetSource, error) { - if result, err := m.DO.Last(); err != nil { - return nil, err - } else { - return result.(*model.ModelDatasetSource), nil - } -} - -func (m modelDatasetSourceDo) Find() ([]*model.ModelDatasetSource, error) { - result, err := m.DO.Find() - return result.([]*model.ModelDatasetSource), err -} - -func (m modelDatasetSourceDo) FindInBatch(batchSize int, fc func(tx gen.Dao, batch int) error) (results []*model.ModelDatasetSource, err error) { - buf := make([]*model.ModelDatasetSource, 0, batchSize) - err = m.DO.FindInBatches(&buf, batchSize, func(tx gen.Dao, batch int) error { - defer func() { results = append(results, buf...) }() - return fc(tx, batch) - }) - return results, err -} - -func (m modelDatasetSourceDo) FindInBatches(result *[]*model.ModelDatasetSource, batchSize int, fc func(tx gen.Dao, batch int) error) error { - return m.DO.FindInBatches(result, batchSize, fc) -} - -func (m modelDatasetSourceDo) Attrs(attrs ...field.AssignExpr) IModelDatasetSourceDo { - return m.withDO(m.DO.Attrs(attrs...)) -} - -func (m modelDatasetSourceDo) Assign(attrs ...field.AssignExpr) IModelDatasetSourceDo { - return m.withDO(m.DO.Assign(attrs...)) -} - -func (m modelDatasetSourceDo) Joins(fields ...field.RelationField) IModelDatasetSourceDo { - for _, _f := range fields { - m = *m.withDO(m.DO.Joins(_f)) - } - return &m -} - -func (m modelDatasetSourceDo) Preload(fields ...field.RelationField) IModelDatasetSourceDo { - for _, _f := range fields { - m = *m.withDO(m.DO.Preload(_f)) - } - return &m -} - -func (m modelDatasetSourceDo) FirstOrInit() (*model.ModelDatasetSource, error) { - if result, err := m.DO.FirstOrInit(); err != nil { - return nil, err - } else { - return result.(*model.ModelDatasetSource), nil - } -} - -func (m modelDatasetSourceDo) FirstOrCreate() (*model.ModelDatasetSource, error) { - if result, err := m.DO.FirstOrCreate(); err != nil { - return nil, err - } else { - return result.(*model.ModelDatasetSource), nil - } -} - -func (m modelDatasetSourceDo) FindByPage(offset int, limit int) (result []*model.ModelDatasetSource, count int64, err error) { - result, err = m.Offset(offset).Limit(limit).Find() - if err != nil { - return - } - - if size := len(result); 0 < limit && 0 < size && size < limit { - count = int64(size + offset) - return - } - - count, err = m.Offset(-1).Limit(-1).Count() - return -} - -func (m modelDatasetSourceDo) ScanByPage(result interface{}, offset int, limit int) (count int64, err error) { - count, err = m.Count() - if err != nil { - return - } - - err = m.Offset(offset).Limit(limit).Scan(result) - return -} - -func (m modelDatasetSourceDo) Scan(result interface{}) (err error) { - return m.DO.Scan(result) -} - -func (m modelDatasetSourceDo) Delete(models ...*model.ModelDatasetSource) (result gen.ResultInfo, err error) { - return m.DO.Delete(models) -} - -func (m *modelDatasetSourceDo) withDO(do gen.Dao) *modelDatasetSourceDo { - m.DO = *do.(*gen.DO) - return m -} diff --git a/backend/dao/query/model_download_submissions.gen.go b/backend/dao/query/model_download_submissions.gen.go deleted file mode 100644 index 5cc3a8197..000000000 --- a/backend/dao/query/model_download_submissions.gen.go +++ /dev/null @@ -1,419 +0,0 @@ -// Code generated by gorm.io/gen. DO NOT EDIT. -// Code generated by gorm.io/gen. DO NOT EDIT. -// Code generated by gorm.io/gen. DO NOT EDIT. - -package query - -import ( - "context" - "database/sql" - - "gorm.io/gorm" - "gorm.io/gorm/clause" - "gorm.io/gorm/schema" - - "gorm.io/gen" - "gorm.io/gen/field" - - "gorm.io/plugin/dbresolver" - - "github.com/raids-lab/crater/dao/model" -) - -func newModelDownloadSubmission(db *gorm.DB, opts ...gen.DOOption) modelDownloadSubmission { - _modelDownloadSubmission := modelDownloadSubmission{} - - _modelDownloadSubmission.modelDownloadSubmissionDo.UseDB(db, opts...) - _modelDownloadSubmission.modelDownloadSubmissionDo.UseModel(&model.ModelDownloadSubmission{}) - - tableName := _modelDownloadSubmission.modelDownloadSubmissionDo.TableName() - _modelDownloadSubmission.ALL = field.NewAsterisk(tableName) - _modelDownloadSubmission.ID = field.NewUint(tableName, "id") - _modelDownloadSubmission.UserID = field.NewUint(tableName, "user_id") - _modelDownloadSubmission.ModelDownloadID = field.NewUint(tableName, "model_download_id") - _modelDownloadSubmission.Action = field.NewString(tableName, "action") - _modelDownloadSubmission.Status = field.NewString(tableName, "status") - _modelDownloadSubmission.CreatedAt = field.NewTime(tableName, "created_at") - _modelDownloadSubmission.CompletedAt = field.NewTime(tableName, "completed_at") - - _modelDownloadSubmission.fillFieldMap() - - return _modelDownloadSubmission -} - -type modelDownloadSubmission struct { - modelDownloadSubmissionDo modelDownloadSubmissionDo - - ALL field.Asterisk - ID field.Uint - UserID field.Uint - ModelDownloadID field.Uint - Action field.String - Status field.String - CreatedAt field.Time - CompletedAt field.Time - - fieldMap map[string]field.Expr -} - -func (m modelDownloadSubmission) Table(newTableName string) *modelDownloadSubmission { - m.modelDownloadSubmissionDo.UseTable(newTableName) - return m.updateTableName(newTableName) -} - -func (m modelDownloadSubmission) As(alias string) *modelDownloadSubmission { - m.modelDownloadSubmissionDo.DO = *(m.modelDownloadSubmissionDo.As(alias).(*gen.DO)) - return m.updateTableName(alias) -} - -func (m *modelDownloadSubmission) updateTableName(table string) *modelDownloadSubmission { - m.ALL = field.NewAsterisk(table) - m.ID = field.NewUint(table, "id") - m.UserID = field.NewUint(table, "user_id") - m.ModelDownloadID = field.NewUint(table, "model_download_id") - m.Action = field.NewString(table, "action") - m.Status = field.NewString(table, "status") - m.CreatedAt = field.NewTime(table, "created_at") - m.CompletedAt = field.NewTime(table, "completed_at") - - m.fillFieldMap() - - return m -} - -func (m *modelDownloadSubmission) WithContext(ctx context.Context) IModelDownloadSubmissionDo { - return m.modelDownloadSubmissionDo.WithContext(ctx) -} - -func (m modelDownloadSubmission) TableName() string { return m.modelDownloadSubmissionDo.TableName() } - -func (m modelDownloadSubmission) Alias() string { return m.modelDownloadSubmissionDo.Alias() } - -func (m modelDownloadSubmission) Columns(cols ...field.Expr) gen.Columns { - return m.modelDownloadSubmissionDo.Columns(cols...) -} - -func (m *modelDownloadSubmission) GetFieldByName(fieldName string) (field.OrderExpr, bool) { - _f, ok := m.fieldMap[fieldName] - if !ok || _f == nil { - return nil, false - } - _oe, ok := _f.(field.OrderExpr) - return _oe, ok -} - -func (m *modelDownloadSubmission) fillFieldMap() { - m.fieldMap = make(map[string]field.Expr, 7) - m.fieldMap["id"] = m.ID - m.fieldMap["user_id"] = m.UserID - m.fieldMap["model_download_id"] = m.ModelDownloadID - m.fieldMap["action"] = m.Action - m.fieldMap["status"] = m.Status - m.fieldMap["created_at"] = m.CreatedAt - m.fieldMap["completed_at"] = m.CompletedAt -} - -func (m modelDownloadSubmission) clone(db *gorm.DB) modelDownloadSubmission { - m.modelDownloadSubmissionDo.ReplaceConnPool(db.Statement.ConnPool) - return m -} - -func (m modelDownloadSubmission) replaceDB(db *gorm.DB) modelDownloadSubmission { - m.modelDownloadSubmissionDo.ReplaceDB(db) - return m -} - -type modelDownloadSubmissionDo struct{ gen.DO } - -type IModelDownloadSubmissionDo interface { - gen.SubQuery - Debug() IModelDownloadSubmissionDo - WithContext(ctx context.Context) IModelDownloadSubmissionDo - WithResult(fc func(tx gen.Dao)) gen.ResultInfo - ReplaceDB(db *gorm.DB) - ReadDB() IModelDownloadSubmissionDo - WriteDB() IModelDownloadSubmissionDo - As(alias string) gen.Dao - Session(config *gorm.Session) IModelDownloadSubmissionDo - Columns(cols ...field.Expr) gen.Columns - Clauses(conds ...clause.Expression) IModelDownloadSubmissionDo - Not(conds ...gen.Condition) IModelDownloadSubmissionDo - Or(conds ...gen.Condition) IModelDownloadSubmissionDo - Select(conds ...field.Expr) IModelDownloadSubmissionDo - Where(conds ...gen.Condition) IModelDownloadSubmissionDo - Order(conds ...field.Expr) IModelDownloadSubmissionDo - Distinct(cols ...field.Expr) IModelDownloadSubmissionDo - Omit(cols ...field.Expr) IModelDownloadSubmissionDo - Join(table schema.Tabler, on ...field.Expr) IModelDownloadSubmissionDo - LeftJoin(table schema.Tabler, on ...field.Expr) IModelDownloadSubmissionDo - RightJoin(table schema.Tabler, on ...field.Expr) IModelDownloadSubmissionDo - Group(cols ...field.Expr) IModelDownloadSubmissionDo - Having(conds ...gen.Condition) IModelDownloadSubmissionDo - Limit(limit int) IModelDownloadSubmissionDo - Offset(offset int) IModelDownloadSubmissionDo - Count() (count int64, err error) - Scopes(funcs ...func(gen.Dao) gen.Dao) IModelDownloadSubmissionDo - Unscoped() IModelDownloadSubmissionDo - Create(values ...*model.ModelDownloadSubmission) error - CreateInBatches(values []*model.ModelDownloadSubmission, batchSize int) error - Save(values ...*model.ModelDownloadSubmission) error - First() (*model.ModelDownloadSubmission, error) - Take() (*model.ModelDownloadSubmission, error) - Last() (*model.ModelDownloadSubmission, error) - Find() ([]*model.ModelDownloadSubmission, error) - FindInBatch(batchSize int, fc func(tx gen.Dao, batch int) error) (results []*model.ModelDownloadSubmission, err error) - FindInBatches(result *[]*model.ModelDownloadSubmission, batchSize int, fc func(tx gen.Dao, batch int) error) error - Pluck(column field.Expr, dest interface{}) error - Delete(...*model.ModelDownloadSubmission) (info gen.ResultInfo, err error) - Update(column field.Expr, value interface{}) (info gen.ResultInfo, err error) - UpdateSimple(columns ...field.AssignExpr) (info gen.ResultInfo, err error) - Updates(value interface{}) (info gen.ResultInfo, err error) - UpdateColumn(column field.Expr, value interface{}) (info gen.ResultInfo, err error) - UpdateColumnSimple(columns ...field.AssignExpr) (info gen.ResultInfo, err error) - UpdateColumns(value interface{}) (info gen.ResultInfo, err error) - UpdateFrom(q gen.SubQuery) gen.Dao - Attrs(attrs ...field.AssignExpr) IModelDownloadSubmissionDo - Assign(attrs ...field.AssignExpr) IModelDownloadSubmissionDo - Joins(fields ...field.RelationField) IModelDownloadSubmissionDo - Preload(fields ...field.RelationField) IModelDownloadSubmissionDo - FirstOrInit() (*model.ModelDownloadSubmission, error) - FirstOrCreate() (*model.ModelDownloadSubmission, error) - FindByPage(offset int, limit int) (result []*model.ModelDownloadSubmission, count int64, err error) - ScanByPage(result interface{}, offset int, limit int) (count int64, err error) - Rows() (*sql.Rows, error) - Row() *sql.Row - Scan(result interface{}) (err error) - Returning(value interface{}, columns ...string) IModelDownloadSubmissionDo - UnderlyingDB() *gorm.DB - schema.Tabler -} - -func (m modelDownloadSubmissionDo) Debug() IModelDownloadSubmissionDo { - return m.withDO(m.DO.Debug()) -} - -func (m modelDownloadSubmissionDo) WithContext(ctx context.Context) IModelDownloadSubmissionDo { - return m.withDO(m.DO.WithContext(ctx)) -} - -func (m modelDownloadSubmissionDo) ReadDB() IModelDownloadSubmissionDo { - return m.Clauses(dbresolver.Read) -} - -func (m modelDownloadSubmissionDo) WriteDB() IModelDownloadSubmissionDo { - return m.Clauses(dbresolver.Write) -} - -func (m modelDownloadSubmissionDo) Session(config *gorm.Session) IModelDownloadSubmissionDo { - return m.withDO(m.DO.Session(config)) -} - -func (m modelDownloadSubmissionDo) Clauses(conds ...clause.Expression) IModelDownloadSubmissionDo { - return m.withDO(m.DO.Clauses(conds...)) -} - -func (m modelDownloadSubmissionDo) Returning(value interface{}, columns ...string) IModelDownloadSubmissionDo { - return m.withDO(m.DO.Returning(value, columns...)) -} - -func (m modelDownloadSubmissionDo) Not(conds ...gen.Condition) IModelDownloadSubmissionDo { - return m.withDO(m.DO.Not(conds...)) -} - -func (m modelDownloadSubmissionDo) Or(conds ...gen.Condition) IModelDownloadSubmissionDo { - return m.withDO(m.DO.Or(conds...)) -} - -func (m modelDownloadSubmissionDo) Select(conds ...field.Expr) IModelDownloadSubmissionDo { - return m.withDO(m.DO.Select(conds...)) -} - -func (m modelDownloadSubmissionDo) Where(conds ...gen.Condition) IModelDownloadSubmissionDo { - return m.withDO(m.DO.Where(conds...)) -} - -func (m modelDownloadSubmissionDo) Order(conds ...field.Expr) IModelDownloadSubmissionDo { - return m.withDO(m.DO.Order(conds...)) -} - -func (m modelDownloadSubmissionDo) Distinct(cols ...field.Expr) IModelDownloadSubmissionDo { - return m.withDO(m.DO.Distinct(cols...)) -} - -func (m modelDownloadSubmissionDo) Omit(cols ...field.Expr) IModelDownloadSubmissionDo { - return m.withDO(m.DO.Omit(cols...)) -} - -func (m modelDownloadSubmissionDo) Join(table schema.Tabler, on ...field.Expr) IModelDownloadSubmissionDo { - return m.withDO(m.DO.Join(table, on...)) -} - -func (m modelDownloadSubmissionDo) LeftJoin(table schema.Tabler, on ...field.Expr) IModelDownloadSubmissionDo { - return m.withDO(m.DO.LeftJoin(table, on...)) -} - -func (m modelDownloadSubmissionDo) RightJoin(table schema.Tabler, on ...field.Expr) IModelDownloadSubmissionDo { - return m.withDO(m.DO.RightJoin(table, on...)) -} - -func (m modelDownloadSubmissionDo) Group(cols ...field.Expr) IModelDownloadSubmissionDo { - return m.withDO(m.DO.Group(cols...)) -} - -func (m modelDownloadSubmissionDo) Having(conds ...gen.Condition) IModelDownloadSubmissionDo { - return m.withDO(m.DO.Having(conds...)) -} - -func (m modelDownloadSubmissionDo) Limit(limit int) IModelDownloadSubmissionDo { - return m.withDO(m.DO.Limit(limit)) -} - -func (m modelDownloadSubmissionDo) Offset(offset int) IModelDownloadSubmissionDo { - return m.withDO(m.DO.Offset(offset)) -} - -func (m modelDownloadSubmissionDo) Scopes(funcs ...func(gen.Dao) gen.Dao) IModelDownloadSubmissionDo { - return m.withDO(m.DO.Scopes(funcs...)) -} - -func (m modelDownloadSubmissionDo) Unscoped() IModelDownloadSubmissionDo { - return m.withDO(m.DO.Unscoped()) -} - -func (m modelDownloadSubmissionDo) Create(values ...*model.ModelDownloadSubmission) error { - if len(values) == 0 { - return nil - } - return m.DO.Create(values) -} - -func (m modelDownloadSubmissionDo) CreateInBatches(values []*model.ModelDownloadSubmission, batchSize int) error { - return m.DO.CreateInBatches(values, batchSize) -} - -// Save : !!! underlying implementation is different with GORM -// The method is equivalent to executing the statement: db.Clauses(clause.OnConflict{UpdateAll: true}).Create(values) -func (m modelDownloadSubmissionDo) Save(values ...*model.ModelDownloadSubmission) error { - if len(values) == 0 { - return nil - } - return m.DO.Save(values) -} - -func (m modelDownloadSubmissionDo) First() (*model.ModelDownloadSubmission, error) { - if result, err := m.DO.First(); err != nil { - return nil, err - } else { - return result.(*model.ModelDownloadSubmission), nil - } -} - -func (m modelDownloadSubmissionDo) Take() (*model.ModelDownloadSubmission, error) { - if result, err := m.DO.Take(); err != nil { - return nil, err - } else { - return result.(*model.ModelDownloadSubmission), nil - } -} - -func (m modelDownloadSubmissionDo) Last() (*model.ModelDownloadSubmission, error) { - if result, err := m.DO.Last(); err != nil { - return nil, err - } else { - return result.(*model.ModelDownloadSubmission), nil - } -} - -func (m modelDownloadSubmissionDo) Find() ([]*model.ModelDownloadSubmission, error) { - result, err := m.DO.Find() - return result.([]*model.ModelDownloadSubmission), err -} - -func (m modelDownloadSubmissionDo) FindInBatch(batchSize int, fc func(tx gen.Dao, batch int) error) (results []*model.ModelDownloadSubmission, err error) { - buf := make([]*model.ModelDownloadSubmission, 0, batchSize) - err = m.DO.FindInBatches(&buf, batchSize, func(tx gen.Dao, batch int) error { - defer func() { results = append(results, buf...) }() - return fc(tx, batch) - }) - return results, err -} - -func (m modelDownloadSubmissionDo) FindInBatches(result *[]*model.ModelDownloadSubmission, batchSize int, fc func(tx gen.Dao, batch int) error) error { - return m.DO.FindInBatches(result, batchSize, fc) -} - -func (m modelDownloadSubmissionDo) Attrs(attrs ...field.AssignExpr) IModelDownloadSubmissionDo { - return m.withDO(m.DO.Attrs(attrs...)) -} - -func (m modelDownloadSubmissionDo) Assign(attrs ...field.AssignExpr) IModelDownloadSubmissionDo { - return m.withDO(m.DO.Assign(attrs...)) -} - -func (m modelDownloadSubmissionDo) Joins(fields ...field.RelationField) IModelDownloadSubmissionDo { - for _, _f := range fields { - m = *m.withDO(m.DO.Joins(_f)) - } - return &m -} - -func (m modelDownloadSubmissionDo) Preload(fields ...field.RelationField) IModelDownloadSubmissionDo { - for _, _f := range fields { - m = *m.withDO(m.DO.Preload(_f)) - } - return &m -} - -func (m modelDownloadSubmissionDo) FirstOrInit() (*model.ModelDownloadSubmission, error) { - if result, err := m.DO.FirstOrInit(); err != nil { - return nil, err - } else { - return result.(*model.ModelDownloadSubmission), nil - } -} - -func (m modelDownloadSubmissionDo) FirstOrCreate() (*model.ModelDownloadSubmission, error) { - if result, err := m.DO.FirstOrCreate(); err != nil { - return nil, err - } else { - return result.(*model.ModelDownloadSubmission), nil - } -} - -func (m modelDownloadSubmissionDo) FindByPage(offset int, limit int) (result []*model.ModelDownloadSubmission, count int64, err error) { - result, err = m.Offset(offset).Limit(limit).Find() - if err != nil { - return - } - - if size := len(result); 0 < limit && 0 < size && size < limit { - count = int64(size + offset) - return - } - - count, err = m.Offset(-1).Limit(-1).Count() - return -} - -func (m modelDownloadSubmissionDo) ScanByPage(result interface{}, offset int, limit int) (count int64, err error) { - count, err = m.Count() - if err != nil { - return - } - - err = m.Offset(offset).Limit(limit).Scan(result) - return -} - -func (m modelDownloadSubmissionDo) Scan(result interface{}) (err error) { - return m.DO.Scan(result) -} - -func (m modelDownloadSubmissionDo) Delete(models ...*model.ModelDownloadSubmission) (result gen.ResultInfo, err error) { - return m.DO.Delete(models) -} - -func (m *modelDownloadSubmissionDo) withDO(do gen.Dao) *modelDownloadSubmissionDo { - m.DO = *do.(*gen.DO) - return m -} diff --git a/backend/dao/query/model_downloads.gen.go b/backend/dao/query/model_downloads.gen.go index 80f022045..26e19d776 100644 --- a/backend/dao/query/model_downloads.gen.go +++ b/backend/dao/query/model_downloads.gen.go @@ -63,16 +63,9 @@ func newModelDownload(db *gorm.DB, opts ...gen.DOOption) modelDownload { _modelDownload.SourceCreatedAt = field.NewTime(tableName, "source_created_at") _modelDownload.SourceUpdatedAt = field.NewTime(tableName, "source_updated_at") _modelDownload.MetadataRefreshedAt = field.NewTime(tableName, "metadata_refreshed_at") - _modelDownload.ModelDatasetSourceID = field.NewUint(tableName, "model_dataset_source_id") _modelDownload.JobName = field.NewString(tableName, "job_name") _modelDownload.CreatorID = field.NewUint(tableName, "creator_id") _modelDownload.ReferenceCount = field.NewInt(tableName, "reference_count") - _modelDownload.ModelDatasetSource = modelDownloadBelongsToModelDatasetSource{ - db: db.Session(&gorm.Session{}), - - RelationField: field.NewRelation("ModelDatasetSource", "model.ModelDatasetSource"), - } - _modelDownload.Creator = modelDownloadBelongsToCreator{ db: db.Session(&gorm.Session{}), @@ -97,49 +90,46 @@ func newModelDownload(db *gorm.DB, opts ...gen.DOOption) modelDownload { type modelDownload struct { modelDownloadDo modelDownloadDo - ALL field.Asterisk - ID field.Uint - CreatedAt field.Time - UpdatedAt field.Time - DeletedAt field.Field - Name field.String - Source field.String - Category field.String - Revision field.String // 版本/分支/commit - Path field.String // 实际下载路径 - SizeBytes field.Int64 // 文件总大小(字节) - DownloadedBytes field.Int64 // 已下载大小(字节) - DownloadSpeed field.String // 下载速度(如: 10MB/s) - Status field.String // 下载状态 - Message field.String // 状态消息(错误信息等) - Logs field.String // 终态时保存的Pod日志(K8s Job被GC后仍可查看) - LogsSavedAt field.Time // 日志保存时间 - Organization field.String // 源站组织或作者 - LogoURL field.String // 源站组织头像地址 - SourceURL field.String // 源站仓库详情地址 - DisplayName field.String // 源站展示名称 - SourceDescription field.String // 源站简介摘要 - SourceReadme field.String // 源站README内容(截断保存) - License field.String // 源站许可证 - Task field.String // 源站任务分类 - Library field.String // 源站框架或库 - ModelType field.String // 源站模型类型 - ParameterCount field.Int64 // 模型参数量 - SourcePrivate field.Bool // 源站是否私有 - SourceGated field.Bool // 源站是否需要申请访问 - SourceLoginRequired field.Bool // 源站是否要求登录下载 - SourceDownloads field.Int64 // 源站下载次数 - SourceLikes field.Int64 // 源站点赞次数 - SourceCreatedAt field.Time // 源站创建时间 - SourceUpdatedAt field.Time // 源站更新时间 - MetadataRefreshedAt field.Time // 源站元数据刷新时间 - ModelDatasetSourceID field.Uint // 模型或数据集外部来源ID - JobName field.String // K8s Job名称 - CreatorID field.Uint // 首个发起下载的用户ID - ReferenceCount field.Int // 提交下载需求的用户计数 - ModelDatasetSource modelDownloadBelongsToModelDatasetSource - - Creator modelDownloadBelongsToCreator + ALL field.Asterisk + ID field.Uint + CreatedAt field.Time + UpdatedAt field.Time + DeletedAt field.Field + Name field.String + Source field.String + Category field.String + Revision field.String // 版本/分支/commit + Path field.String // 实际下载路径 + SizeBytes field.Int64 // 文件总大小(字节) + DownloadedBytes field.Int64 // 已下载大小(字节) + DownloadSpeed field.String // 下载速度(如: 10MB/s) + Status field.String // 下载状态 + Message field.String // 状态消息(错误信息等) + Logs field.String // 终态时保存的Pod日志(K8s Job被GC后仍可查看) + LogsSavedAt field.Time // 日志保存时间 + Organization field.String // 源站组织或作者 + LogoURL field.String // 源站组织头像地址 + SourceURL field.String // 源站仓库详情地址 + DisplayName field.String // 源站展示名称 + SourceDescription field.String // 源站简介摘要 + SourceReadme field.String // 源站README内容(截断保存) + License field.String // 源站许可证 + Task field.String // 源站任务分类 + Library field.String // 源站框架或库 + ModelType field.String // 源站模型类型 + ParameterCount field.Int64 // 模型参数量 + SourcePrivate field.Bool // 源站是否私有 + SourceGated field.Bool // 源站是否需要申请访问 + SourceLoginRequired field.Bool // 源站是否要求登录下载 + SourceDownloads field.Int64 // 源站下载次数 + SourceLikes field.Int64 // 源站点赞次数 + SourceCreatedAt field.Time // 源站创建时间 + SourceUpdatedAt field.Time // 源站更新时间 + MetadataRefreshedAt field.Time // 源站元数据刷新时间 + JobName field.String // K8s Job名称 + CreatorID field.Uint // 首个发起下载的用户ID + ReferenceCount field.Int // 引用计数 + Creator modelDownloadBelongsToCreator fieldMap map[string]field.Expr } @@ -191,7 +181,6 @@ func (m *modelDownload) updateTableName(table string) *modelDownload { m.SourceCreatedAt = field.NewTime(table, "source_created_at") m.SourceUpdatedAt = field.NewTime(table, "source_updated_at") m.MetadataRefreshedAt = field.NewTime(table, "metadata_refreshed_at") - m.ModelDatasetSourceID = field.NewUint(table, "model_dataset_source_id") m.JobName = field.NewString(table, "job_name") m.CreatorID = field.NewUint(table, "creator_id") m.ReferenceCount = field.NewInt(table, "reference_count") @@ -223,7 +212,7 @@ func (m *modelDownload) GetFieldByName(fieldName string) (field.OrderExpr, bool) } func (m *modelDownload) fillFieldMap() { - m.fieldMap = make(map[string]field.Expr, 41) + m.fieldMap = make(map[string]field.Expr, 39) m.fieldMap["id"] = m.ID m.fieldMap["created_at"] = m.CreatedAt m.fieldMap["updated_at"] = m.UpdatedAt @@ -259,7 +248,6 @@ func (m *modelDownload) fillFieldMap() { m.fieldMap["source_created_at"] = m.SourceCreatedAt m.fieldMap["source_updated_at"] = m.SourceUpdatedAt m.fieldMap["metadata_refreshed_at"] = m.MetadataRefreshedAt - m.fieldMap["model_dataset_source_id"] = m.ModelDatasetSourceID m.fieldMap["job_name"] = m.JobName m.fieldMap["creator_id"] = m.CreatorID m.fieldMap["reference_count"] = m.ReferenceCount @@ -268,8 +256,6 @@ func (m *modelDownload) fillFieldMap() { func (m modelDownload) clone(db *gorm.DB) modelDownload { m.modelDownloadDo.ReplaceConnPool(db.Statement.ConnPool) - m.ModelDatasetSource.db = db.Session(&gorm.Session{Initialized: true}) - m.ModelDatasetSource.db.Statement.ConnPool = db.Statement.ConnPool m.Creator.db = db.Session(&gorm.Session{Initialized: true}) m.Creator.db.Statement.ConnPool = db.Statement.ConnPool return m @@ -277,92 +263,10 @@ func (m modelDownload) clone(db *gorm.DB) modelDownload { func (m modelDownload) replaceDB(db *gorm.DB) modelDownload { m.modelDownloadDo.ReplaceDB(db) - m.ModelDatasetSource.db = db.Session(&gorm.Session{}) m.Creator.db = db.Session(&gorm.Session{}) return m } -type modelDownloadBelongsToModelDatasetSource struct { - db *gorm.DB - - field.RelationField -} - -func (a modelDownloadBelongsToModelDatasetSource) Where(conds ...field.Expr) *modelDownloadBelongsToModelDatasetSource { - if len(conds) == 0 { - return &a - } - - exprs := make([]clause.Expression, 0, len(conds)) - for _, cond := range conds { - exprs = append(exprs, cond.BeCond().(clause.Expression)) - } - a.db = a.db.Clauses(clause.Where{Exprs: exprs}) - return &a -} - -func (a modelDownloadBelongsToModelDatasetSource) WithContext(ctx context.Context) *modelDownloadBelongsToModelDatasetSource { - a.db = a.db.WithContext(ctx) - return &a -} - -func (a modelDownloadBelongsToModelDatasetSource) Session(session *gorm.Session) *modelDownloadBelongsToModelDatasetSource { - a.db = a.db.Session(session) - return &a -} - -func (a modelDownloadBelongsToModelDatasetSource) Model(m *model.ModelDownload) *modelDownloadBelongsToModelDatasetSourceTx { - return &modelDownloadBelongsToModelDatasetSourceTx{a.db.Model(m).Association(a.Name())} -} - -func (a modelDownloadBelongsToModelDatasetSource) Unscoped() *modelDownloadBelongsToModelDatasetSource { - a.db = a.db.Unscoped() - return &a -} - -type modelDownloadBelongsToModelDatasetSourceTx struct{ tx *gorm.Association } - -func (a modelDownloadBelongsToModelDatasetSourceTx) Find() (result *model.ModelDatasetSource, err error) { - return result, a.tx.Find(&result) -} - -func (a modelDownloadBelongsToModelDatasetSourceTx) Append(values ...*model.ModelDatasetSource) (err error) { - targetValues := make([]interface{}, len(values)) - for i, v := range values { - targetValues[i] = v - } - return a.tx.Append(targetValues...) -} - -func (a modelDownloadBelongsToModelDatasetSourceTx) Replace(values ...*model.ModelDatasetSource) (err error) { - targetValues := make([]interface{}, len(values)) - for i, v := range values { - targetValues[i] = v - } - return a.tx.Replace(targetValues...) -} - -func (a modelDownloadBelongsToModelDatasetSourceTx) Delete(values ...*model.ModelDatasetSource) (err error) { - targetValues := make([]interface{}, len(values)) - for i, v := range values { - targetValues[i] = v - } - return a.tx.Delete(targetValues...) -} - -func (a modelDownloadBelongsToModelDatasetSourceTx) Clear() error { - return a.tx.Clear() -} - -func (a modelDownloadBelongsToModelDatasetSourceTx) Count() int64 { - return a.tx.Count() -} - -func (a modelDownloadBelongsToModelDatasetSourceTx) Unscoped() *modelDownloadBelongsToModelDatasetSourceTx { - a.tx = a.tx.Unscoped() - return &a -} - type modelDownloadBelongsToCreator struct { db *gorm.DB diff --git a/backend/dao/query/resources.gen.go b/backend/dao/query/resources.gen.go index 9707b5198..7d21414a1 100644 --- a/backend/dao/query/resources.gen.go +++ b/backend/dao/query/resources.gen.go @@ -74,7 +74,7 @@ type resource struct { Format field.String // 资源格式 Priority field.Int // 优先级 Label field.String // 用于显示的别名 - UnitPrice field.Int64 // 资源单位价格(内部微点, 展示为点数/单位/小时) + UnitPrice field.Int64 // 资源单位价格(点数/单位/分钟) Type field.String // 资源类型 Networks resourceManyToManyNetworks diff --git a/backend/dao/query/user_accounts.gen.go b/backend/dao/query/user_accounts.gen.go index 9f0dc144e..20ef67dce 100644 --- a/backend/dao/query/user_accounts.gen.go +++ b/backend/dao/query/user_accounts.gen.go @@ -58,8 +58,8 @@ type userAccount struct { Role field.Uint8 // 用户在账户中的角色 (user, admin) AccessMode field.Uint8 // 用户在账户空间的访问模式 (na, ro, rw) Quota field.Field // 用户在账户中的资源配额 - BillingIssueAmountOverride field.Int64 // 用户在账户内的周期发放额度覆盖(内部微点, 为空表示沿用账户配置) - PeriodFreeBalance field.Int64 // 用户在当前周期的免费额度剩余(内部微点) + BillingIssueAmountOverride field.Int64 // 用户在账户内的周期发放额度覆盖(为空表示沿用账户配置) + PeriodFreeBalance field.Int64 // 用户在当前周期的免费额度剩余 fieldMap map[string]field.Expr } diff --git a/backend/dao/query/user_model_downloads.gen.go b/backend/dao/query/user_model_downloads.gen.go index 91933d5dd..aab089a8c 100644 --- a/backend/dao/query/user_model_downloads.gen.go +++ b/backend/dao/query/user_model_downloads.gen.go @@ -52,11 +52,6 @@ func newUserModelDownload(db *gorm.DB, opts ...gen.DOOption) userModelDownload { db: db.Session(&gorm.Session{}), RelationField: field.NewRelation("ModelDownload", "model.ModelDownload"), - ModelDatasetSource: struct { - field.RelationField - }{ - RelationField: field.NewRelation("ModelDownload.ModelDatasetSource", "model.ModelDatasetSource"), - }, Creator: struct { field.RelationField }{ @@ -245,9 +240,6 @@ type userModelDownloadBelongsToModelDownload struct { field.RelationField - ModelDatasetSource struct { - field.RelationField - } Creator struct { field.RelationField } diff --git a/backend/dao/query/users.gen.go b/backend/dao/query/users.gen.go index f63c1c78e..98aefdbc4 100644 --- a/backend/dao/query/users.gen.go +++ b/backend/dao/query/users.gen.go @@ -74,7 +74,7 @@ type user struct { Status field.Uint8 // 用户状态 (pending, active, inactive) Space field.String // 用户空间绝对路径 ImageQuota field.Int64 // 用户在镜像仓库的配额 - ExtraBalance field.Int64 // 用户额外点数余额(内部微点, 充值/奖励) + ExtraBalance field.Int64 // 用户额外点数余额(充值/奖励) LastEmailVerifiedAt field.Time // 最后一次邮箱验证时间 Attributes field.Field // 用户的额外属性 (昵称、邮箱、电话、头像等) UserAccounts userHasManyUserAccounts diff --git a/backend/docs/docs.go b/backend/docs/docs.go index 0440209f8..6749d9b2b 100644 --- a/backend/docs/docs.go +++ b/backend/docs/docs.go @@ -184,44 +184,6 @@ const docTemplate = `{ } } }, - "/dataset/source-logo/{sourceId}": { - "get": { - "description": "返回平台缓存的来源 Logo，不要求浏览器访问外部模型站点", - "produces": [ - "image/png", - "image/jpeg", - "image/webp", - "image/svg+xml" - ], - "tags": [ - "Dataset" - ], - "summary": "获取平台缓存的模型或数据集来源 Logo", - "parameters": [ - { - "type": "integer", - "description": "来源 ID", - "name": "sourceId", - "in": "path", - "required": true - } - ], - "responses": { - "200": { - "description": "OK", - "schema": { - "type": "file" - } - }, - "400": { - "description": "Bad Request" - }, - "404": { - "description": "Not Found" - } - } - } - }, "/metrics": { "get": { "security": [ @@ -1512,68 +1474,6 @@ const docTemplate = `{ "responses": {} } }, - "/v1/admin/images/cudabaseimage": { - "post": { - "security": [ - { - "Bearer": [] - } - ], - "description": "添加新的CUDA基础镜像到系统中", - "consumes": [ - "application/json" - ], - "produces": [ - "application/json" - ], - "tags": [ - "ImagePack" - ], - "summary": "添加CUDA基础镜像", - "parameters": [ - { - "description": "CUDA基础镜像信息", - "name": "data", - "in": "body", - "required": true, - "schema": { - "$ref": "#/definitions/internal_handler_image.CudaBaseImageCreateRequest" - } - } - ], - "responses": {} - } - }, - "/v1/admin/images/cudabaseimage/{id}": { - "delete": { - "security": [ - { - "Bearer": [] - } - ], - "description": "根据ID删除指定的CUDA基础镜像", - "consumes": [ - "application/json" - ], - "produces": [ - "application/json" - ], - "tags": [ - "ImagePack" - ], - "summary": "删除CUDA基础镜像", - "parameters": [ - { - "type": "integer", - "description": "CUDA基础镜像ID", - "name": "id", - "in": "path", - "required": true - } - ], - "responses": {} - } - }, "/v1/admin/images/description": { "post": { "security": [ @@ -3380,68 +3280,6 @@ const docTemplate = `{ } } }, - "/v1/admin/system-config/model-download-limit": { - "get": { - "security": [ - { - "Bearer": [] - } - ], - "description": "获取下载额度和白名单用户 ID", - "produces": [ - "application/json" - ], - "tags": [ - "SystemConfig" - ], - "summary": "管理员获取模型与数据集下载额度", - "responses": { - "200": { - "description": "OK", - "schema": { - "$ref": "#/definitions/github_com_raids-lab_crater_internal_resputil.Response-internal_handler_AdminModelDownloadLimitConfigResp" - } - } - } - }, - "put": { - "security": [ - { - "Bearer": [] - } - ], - "description": "配置所有用户的并发任务上限、滚动窗口成功下载上限和豁免白名单", - "consumes": [ - "application/json" - ], - "produces": [ - "application/json" - ], - "tags": [ - "SystemConfig" - ], - "summary": "更新模型与数据集下载额度", - "parameters": [ - { - "description": "下载额度配置", - "name": "data", - "in": "body", - "required": true, - "schema": { - "$ref": "#/definitions/internal_handler.UpdateModelDownloadLimitConfigReq" - } - } - ], - "responses": { - "200": { - "description": "OK", - "schema": { - "$ref": "#/definitions/github_com_raids-lab_crater_internal_resputil.Response-string" - } - } - } - } - }, "/v1/admin/system-config/prequeue": { "get": { "security": [ @@ -5680,6 +5518,66 @@ const docTemplate = `{ ], "summary": "获取所有Cuda基础镜像", "responses": {} + }, + "post": { + "security": [ + { + "Bearer": [] + } + ], + "description": "添加新的CUDA基础镜像到系统中", + "consumes": [ + "application/json" + ], + "produces": [ + "application/json" + ], + "tags": [ + "ImagePack" + ], + "summary": "添加CUDA基础镜像", + "parameters": [ + { + "description": "CUDA基础镜像信息", + "name": "data", + "in": "body", + "required": true, + "schema": { + "$ref": "#/definitions/internal_handler_image.CudaBaseImageCreateRequest" + } + } + ], + "responses": {} + } + }, + "/v1/images/cudabaseimage/{id}": { + "delete": { + "security": [ + { + "Bearer": [] + } + ], + "description": "根据ID删除指定的CUDA基础镜像", + "consumes": [ + "application/json" + ], + "produces": [ + "application/json" + ], + "tags": [ + "ImagePack" + ], + "summary": "删除CUDA基础镜像", + "parameters": [ + { + "type": "integer", + "description": "CUDA基础镜像ID", + "name": "id", + "in": "path", + "required": true + } + ], + "responses": {} } }, "/v1/images/deleteimage": { @@ -6659,7 +6557,7 @@ const docTemplate = `{ "Bearer": [] } ], - "description": "删除下载任务记录(仅平台管理员),已下载的文件保留在存储中", + "description": "删除下载任务记录(仅创建者或管理员),已下载的文件保留在存储中", "consumes": [ "application/json" ], @@ -6835,7 +6733,7 @@ const docTemplate = `{ "required": true }, { - "description": "可选的临时访问令牌和重试版本", + "description": "可选的临时访问令牌", "name": "data", "in": "body", "schema": { @@ -8603,31 +8501,6 @@ const docTemplate = `{ } } }, - "/v1/system-config/model-download-limit": { - "get": { - "security": [ - { - "Bearer": [] - } - ], - "description": "获取当前用户的并发任务上限、滚动窗口成功下载上限和白名单豁免状态", - "produces": [ - "application/json" - ], - "tags": [ - "SystemConfig" - ], - "summary": "获取模型与数据集下载额度", - "responses": { - "200": { - "description": "OK", - "schema": { - "$ref": "#/definitions/github_com_raids-lab_crater_internal_resputil.Response-internal_handler_ModelDownloadLimitConfigResp" - } - } - } - } - }, "/v1/token/verify": { "get": { "security": [ @@ -10858,21 +10731,6 @@ const docTemplate = `{ } } }, - "github_com_raids-lab_crater_internal_resputil.Response-internal_handler_AdminModelDownloadLimitConfigResp": { - "type": "object", - "properties": { - "code": { - "description": "依然保持 int (ErrorCode) 类型", - "type": "integer" - }, - "data": { - "$ref": "#/definitions/internal_handler.AdminModelDownloadLimitConfigResp" - }, - "msg": { - "type": "string" - } - } - }, "github_com_raids-lab_crater_internal_resputil.Response-internal_handler_ApprovalOrderResp": { "type": "object", "properties": { @@ -10993,21 +10851,6 @@ const docTemplate = `{ } } }, - "github_com_raids-lab_crater_internal_resputil.Response-internal_handler_ModelDownloadLimitConfigResp": { - "type": "object", - "properties": { - "code": { - "description": "依然保持 int (ErrorCode) 类型", - "type": "integer" - }, - "data": { - "$ref": "#/definitions/internal_handler.ModelDownloadLimitConfigResp" - }, - "msg": { - "type": "string" - } - } - }, "github_com_raids-lab_crater_internal_resputil.Response-internal_handler_ModelDownloadListResp": { "type": "object", "properties": { @@ -11713,29 +11556,6 @@ const docTemplate = `{ } } }, - "internal_handler.AdminModelDownloadLimitConfigResp": { - "type": "object", - "properties": { - "enabled": { - "type": "boolean" - }, - "maxConcurrent": { - "type": "integer" - }, - "maxSuccessfulDownloads": { - "type": "integer" - }, - "whitelistUserIds": { - "type": "array", - "items": { - "type": "integer" - } - }, - "windowHours": { - "type": "integer" - } - } - }, "internal_handler.ApprovalOrderResp": { "type": "object", "properties": { @@ -11929,10 +11749,6 @@ const docTemplate = `{ "internal_handler.DownloadActionReq": { "type": "object", "properties": { - "revision": { - "description": "Revision is optional and only used by retry. A non-nil empty value means\n\"use the source default branch\" while preserving the failed record's path.", - "type": "string" - }, "token": { "type": "string" } @@ -12218,26 +12034,6 @@ const docTemplate = `{ } } }, - "internal_handler.ModelDownloadLimitConfigResp": { - "type": "object", - "properties": { - "enabled": { - "type": "boolean" - }, - "exempt": { - "type": "boolean" - }, - "maxConcurrent": { - "type": "integer" - }, - "maxSuccessfulDownloads": { - "type": "integer" - }, - "windowHours": { - "type": "integer" - } - } - }, "internal_handler.ModelDownloadListResp": { "type": "object", "properties": { @@ -12262,9 +12058,6 @@ const docTemplate = `{ "internal_handler.ModelDownloadResp": { "type": "object", "properties": { - "canDelete": { - "type": "boolean" - }, "canManage": { "type": "boolean" }, @@ -12317,21 +12110,8 @@ const docTemplate = `{ "type": "string" }, "referenceCount": { - "description": "Deprecated: use requesterCount.", - "type": "integer" - }, - "relation": { - "type": "string" - }, - "requesterCount": { "type": "integer" }, - "requesters": { - "type": "array", - "items": { - "$ref": "#/definitions/github_com_raids-lab_crater_dao_model.UserInfo" - } - }, "revision": { "type": "string" }, @@ -12817,35 +12597,6 @@ const docTemplate = `{ } } }, - "internal_handler.UpdateModelDownloadLimitConfigReq": { - "type": "object", - "required": [ - "enabled", - "maxConcurrent", - "maxSuccessfulDownloads", - "windowHours" - ], - "properties": { - "enabled": { - "type": "boolean" - }, - "maxConcurrent": { - "type": "integer" - }, - "maxSuccessfulDownloads": { - "type": "integer" - }, - "whitelistUserIds": { - "type": "array", - "items": { - "type": "integer" - } - }, - "windowHours": { - "type": "integer" - } - } - }, "internal_handler.UpdatePrequeueConfigReq": { "type": "object", "required": [ diff --git a/backend/docs/swagger.json b/backend/docs/swagger.json index 9c0b319a1..8fcdb2f85 100644 --- a/backend/docs/swagger.json +++ b/backend/docs/swagger.json @@ -176,44 +176,6 @@ } } }, - "/dataset/source-logo/{sourceId}": { - "get": { - "description": "返回平台缓存的来源 Logo，不要求浏览器访问外部模型站点", - "produces": [ - "image/png", - "image/jpeg", - "image/webp", - "image/svg+xml" - ], - "tags": [ - "Dataset" - ], - "summary": "获取平台缓存的模型或数据集来源 Logo", - "parameters": [ - { - "type": "integer", - "description": "来源 ID", - "name": "sourceId", - "in": "path", - "required": true - } - ], - "responses": { - "200": { - "description": "OK", - "schema": { - "type": "file" - } - }, - "400": { - "description": "Bad Request" - }, - "404": { - "description": "Not Found" - } - } - } - }, "/metrics": { "get": { "security": [ @@ -1504,68 +1466,6 @@ "responses": {} } }, - "/v1/admin/images/cudabaseimage": { - "post": { - "security": [ - { - "Bearer": [] - } - ], - "description": "添加新的CUDA基础镜像到系统中", - "consumes": [ - "application/json" - ], - "produces": [ - "application/json" - ], - "tags": [ - "ImagePack" - ], - "summary": "添加CUDA基础镜像", - "parameters": [ - { - "description": "CUDA基础镜像信息", - "name": "data", - "in": "body", - "required": true, - "schema": { - "$ref": "#/definitions/internal_handler_image.CudaBaseImageCreateRequest" - } - } - ], - "responses": {} - } - }, - "/v1/admin/images/cudabaseimage/{id}": { - "delete": { - "security": [ - { - "Bearer": [] - } - ], - "description": "根据ID删除指定的CUDA基础镜像", - "consumes": [ - "application/json" - ], - "produces": [ - "application/json" - ], - "tags": [ - "ImagePack" - ], - "summary": "删除CUDA基础镜像", - "parameters": [ - { - "type": "integer", - "description": "CUDA基础镜像ID", - "name": "id", - "in": "path", - "required": true - } - ], - "responses": {} - } - }, "/v1/admin/images/description": { "post": { "security": [ @@ -3372,68 +3272,6 @@ } } }, - "/v1/admin/system-config/model-download-limit": { - "get": { - "security": [ - { - "Bearer": [] - } - ], - "description": "获取下载额度和白名单用户 ID", - "produces": [ - "application/json" - ], - "tags": [ - "SystemConfig" - ], - "summary": "管理员获取模型与数据集下载额度", - "responses": { - "200": { - "description": "OK", - "schema": { - "$ref": "#/definitions/github_com_raids-lab_crater_internal_resputil.Response-internal_handler_AdminModelDownloadLimitConfigResp" - } - } - } - }, - "put": { - "security": [ - { - "Bearer": [] - } - ], - "description": "配置所有用户的并发任务上限、滚动窗口成功下载上限和豁免白名单", - "consumes": [ - "application/json" - ], - "produces": [ - "application/json" - ], - "tags": [ - "SystemConfig" - ], - "summary": "更新模型与数据集下载额度", - "parameters": [ - { - "description": "下载额度配置", - "name": "data", - "in": "body", - "required": true, - "schema": { - "$ref": "#/definitions/internal_handler.UpdateModelDownloadLimitConfigReq" - } - } - ], - "responses": { - "200": { - "description": "OK", - "schema": { - "$ref": "#/definitions/github_com_raids-lab_crater_internal_resputil.Response-string" - } - } - } - } - }, "/v1/admin/system-config/prequeue": { "get": { "security": [ @@ -5672,6 +5510,66 @@ ], "summary": "获取所有Cuda基础镜像", "responses": {} + }, + "post": { + "security": [ + { + "Bearer": [] + } + ], + "description": "添加新的CUDA基础镜像到系统中", + "consumes": [ + "application/json" + ], + "produces": [ + "application/json" + ], + "tags": [ + "ImagePack" + ], + "summary": "添加CUDA基础镜像", + "parameters": [ + { + "description": "CUDA基础镜像信息", + "name": "data", + "in": "body", + "required": true, + "schema": { + "$ref": "#/definitions/internal_handler_image.CudaBaseImageCreateRequest" + } + } + ], + "responses": {} + } + }, + "/v1/images/cudabaseimage/{id}": { + "delete": { + "security": [ + { + "Bearer": [] + } + ], + "description": "根据ID删除指定的CUDA基础镜像", + "consumes": [ + "application/json" + ], + "produces": [ + "application/json" + ], + "tags": [ + "ImagePack" + ], + "summary": "删除CUDA基础镜像", + "parameters": [ + { + "type": "integer", + "description": "CUDA基础镜像ID", + "name": "id", + "in": "path", + "required": true + } + ], + "responses": {} } }, "/v1/images/deleteimage": { @@ -6651,7 +6549,7 @@ "Bearer": [] } ], - "description": "删除下载任务记录(仅平台管理员),已下载的文件保留在存储中", + "description": "删除下载任务记录(仅创建者或管理员),已下载的文件保留在存储中", "consumes": [ "application/json" ], @@ -6827,7 +6725,7 @@ "required": true }, { - "description": "可选的临时访问令牌和重试版本", + "description": "可选的临时访问令牌", "name": "data", "in": "body", "schema": { @@ -8595,31 +8493,6 @@ } } }, - "/v1/system-config/model-download-limit": { - "get": { - "security": [ - { - "Bearer": [] - } - ], - "description": "获取当前用户的并发任务上限、滚动窗口成功下载上限和白名单豁免状态", - "produces": [ - "application/json" - ], - "tags": [ - "SystemConfig" - ], - "summary": "获取模型与数据集下载额度", - "responses": { - "200": { - "description": "OK", - "schema": { - "$ref": "#/definitions/github_com_raids-lab_crater_internal_resputil.Response-internal_handler_ModelDownloadLimitConfigResp" - } - } - } - } - }, "/v1/token/verify": { "get": { "security": [ @@ -10850,21 +10723,6 @@ } } }, - "github_com_raids-lab_crater_internal_resputil.Response-internal_handler_AdminModelDownloadLimitConfigResp": { - "type": "object", - "properties": { - "code": { - "description": "依然保持 int (ErrorCode) 类型", - "type": "integer" - }, - "data": { - "$ref": "#/definitions/internal_handler.AdminModelDownloadLimitConfigResp" - }, - "msg": { - "type": "string" - } - } - }, "github_com_raids-lab_crater_internal_resputil.Response-internal_handler_ApprovalOrderResp": { "type": "object", "properties": { @@ -10985,21 +10843,6 @@ } } }, - "github_com_raids-lab_crater_internal_resputil.Response-internal_handler_ModelDownloadLimitConfigResp": { - "type": "object", - "properties": { - "code": { - "description": "依然保持 int (ErrorCode) 类型", - "type": "integer" - }, - "data": { - "$ref": "#/definitions/internal_handler.ModelDownloadLimitConfigResp" - }, - "msg": { - "type": "string" - } - } - }, "github_com_raids-lab_crater_internal_resputil.Response-internal_handler_ModelDownloadListResp": { "type": "object", "properties": { @@ -11705,29 +11548,6 @@ } } }, - "internal_handler.AdminModelDownloadLimitConfigResp": { - "type": "object", - "properties": { - "enabled": { - "type": "boolean" - }, - "maxConcurrent": { - "type": "integer" - }, - "maxSuccessfulDownloads": { - "type": "integer" - }, - "whitelistUserIds": { - "type": "array", - "items": { - "type": "integer" - } - }, - "windowHours": { - "type": "integer" - } - } - }, "internal_handler.ApprovalOrderResp": { "type": "object", "properties": { @@ -11921,10 +11741,6 @@ "internal_handler.DownloadActionReq": { "type": "object", "properties": { - "revision": { - "description": "Revision is optional and only used by retry. A non-nil empty value means\n\"use the source default branch\" while preserving the failed record's path.", - "type": "string" - }, "token": { "type": "string" } @@ -12210,26 +12026,6 @@ } } }, - "internal_handler.ModelDownloadLimitConfigResp": { - "type": "object", - "properties": { - "enabled": { - "type": "boolean" - }, - "exempt": { - "type": "boolean" - }, - "maxConcurrent": { - "type": "integer" - }, - "maxSuccessfulDownloads": { - "type": "integer" - }, - "windowHours": { - "type": "integer" - } - } - }, "internal_handler.ModelDownloadListResp": { "type": "object", "properties": { @@ -12254,9 +12050,6 @@ "internal_handler.ModelDownloadResp": { "type": "object", "properties": { - "canDelete": { - "type": "boolean" - }, "canManage": { "type": "boolean" }, @@ -12309,21 +12102,8 @@ "type": "string" }, "referenceCount": { - "description": "Deprecated: use requesterCount.", - "type": "integer" - }, - "relation": { - "type": "string" - }, - "requesterCount": { "type": "integer" }, - "requesters": { - "type": "array", - "items": { - "$ref": "#/definitions/github_com_raids-lab_crater_dao_model.UserInfo" - } - }, "revision": { "type": "string" }, @@ -12809,35 +12589,6 @@ } } }, - "internal_handler.UpdateModelDownloadLimitConfigReq": { - "type": "object", - "required": [ - "enabled", - "maxConcurrent", - "maxSuccessfulDownloads", - "windowHours" - ], - "properties": { - "enabled": { - "type": "boolean" - }, - "maxConcurrent": { - "type": "integer" - }, - "maxSuccessfulDownloads": { - "type": "integer" - }, - "whitelistUserIds": { - "type": "array", - "items": { - "type": "integer" - } - }, - "windowHours": { - "type": "integer" - } - } - }, "internal_handler.UpdatePrequeueConfigReq": { "type": "object", "required": [ diff --git a/backend/docs/swagger.yaml b/backend/docs/swagger.yaml index 4c9e6aa61..53dc439a4 100644 --- a/backend/docs/swagger.yaml +++ b/backend/docs/swagger.yaml @@ -573,16 +573,6 @@ definitions: msg: type: string type: object - github_com_raids-lab_crater_internal_resputil.Response-internal_handler_AdminModelDownloadLimitConfigResp: - properties: - code: - description: 依然保持 int (ErrorCode) 类型 - type: integer - data: - $ref: '#/definitions/internal_handler.AdminModelDownloadLimitConfigResp' - msg: - type: string - type: object github_com_raids-lab_crater_internal_resputil.Response-internal_handler_ApprovalOrderResp: properties: code: @@ -663,16 +653,6 @@ definitions: msg: type: string type: object - github_com_raids-lab_crater_internal_resputil.Response-internal_handler_ModelDownloadLimitConfigResp: - properties: - code: - description: 依然保持 int (ErrorCode) 类型 - type: integer - data: - $ref: '#/definitions/internal_handler.ModelDownloadLimitConfigResp' - msg: - type: string - type: object github_com_raids-lab_crater_internal_resputil.Response-internal_handler_ModelDownloadListResp: properties: code: @@ -1141,21 +1121,6 @@ definitions: username: type: string type: object - internal_handler.AdminModelDownloadLimitConfigResp: - properties: - enabled: - type: boolean - maxConcurrent: - type: integer - maxSuccessfulDownloads: - type: integer - whitelistUserIds: - items: - type: integer - type: array - windowHours: - type: integer - type: object internal_handler.ApprovalOrderResp: properties: content: @@ -1289,11 +1254,6 @@ definitions: type: object internal_handler.DownloadActionReq: properties: - revision: - description: |- - Revision is optional and only used by retry. A non-nil empty value means - "use the source default branch" while preserving the failed record's path. - type: string token: type: string type: object @@ -1482,19 +1442,6 @@ definitions: user: $ref: '#/definitions/github_com_raids-lab_crater_dao_model.UserAttribute' type: object - internal_handler.ModelDownloadLimitConfigResp: - properties: - enabled: - type: boolean - exempt: - type: boolean - maxConcurrent: - type: integer - maxSuccessfulDownloads: - type: integer - windowHours: - type: integer - type: object internal_handler.ModelDownloadListResp: properties: items: @@ -1511,8 +1458,6 @@ definitions: type: object internal_handler.ModelDownloadResp: properties: - canDelete: - type: boolean canManage: type: boolean canViewLogs: @@ -1548,16 +1493,7 @@ definitions: path: type: string referenceCount: - description: 'Deprecated: use requesterCount.' type: integer - relation: - type: string - requesterCount: - type: integer - requesters: - items: - $ref: '#/definitions/github_com_raids-lab_crater_dao_model.UserInfo' - type: array revision: type: string sizeBytes: @@ -1878,26 +1814,6 @@ definitions: - baseUrl - modelName type: object - internal_handler.UpdateModelDownloadLimitConfigReq: - properties: - enabled: - type: boolean - maxConcurrent: - type: integer - maxSuccessfulDownloads: - type: integer - whitelistUserIds: - items: - type: integer - type: array - windowHours: - type: integer - required: - - enabled - - maxConcurrent - - maxSuccessfulDownloads - - windowHours - type: object internal_handler.UpdatePrequeueConfigReq: properties: activateTickerIntervalSeconds: @@ -2776,32 +2692,6 @@ paths: summary: 获取后端用户认证模式 tags: - Auth - /dataset/source-logo/{sourceId}: - get: - description: 返回平台缓存的来源 Logo，不要求浏览器访问外部模型站点 - parameters: - - description: 来源 ID - in: path - name: sourceId - required: true - type: integer - produces: - - image/png - - image/jpeg - - image/webp - - image/svg+xml - responses: - "200": - description: OK - schema: - type: file - "400": - description: Bad Request - "404": - description: Not Found - summary: 获取平台缓存的模型或数据集来源 Logo - tags: - - Dataset /metrics: get: consumes: @@ -3621,45 +3511,6 @@ paths: summary: 管理员更新镜像架构 tags: - ImagePack - /v1/admin/images/cudabaseimage: - post: - consumes: - - application/json - description: 添加新的CUDA基础镜像到系统中 - parameters: - - description: CUDA基础镜像信息 - in: body - name: data - required: true - schema: - $ref: '#/definitions/internal_handler_image.CudaBaseImageCreateRequest' - produces: - - application/json - responses: {} - security: - - Bearer: [] - summary: 添加CUDA基础镜像 - tags: - - ImagePack - /v1/admin/images/cudabaseimage/{id}: - delete: - consumes: - - application/json - description: 根据ID删除指定的CUDA基础镜像 - parameters: - - description: CUDA基础镜像ID - in: path - name: id - required: true - type: integer - produces: - - application/json - responses: {} - security: - - Bearer: [] - summary: 删除CUDA基础镜像 - tags: - - ImagePack /v1/admin/images/description: post: consumes: @@ -4817,44 +4668,6 @@ paths: summary: 更新 LLM 配置 tags: - SystemConfig - /v1/admin/system-config/model-download-limit: - get: - description: 获取下载额度和白名单用户 ID - produces: - - application/json - responses: - "200": - description: OK - schema: - $ref: '#/definitions/github_com_raids-lab_crater_internal_resputil.Response-internal_handler_AdminModelDownloadLimitConfigResp' - security: - - Bearer: [] - summary: 管理员获取模型与数据集下载额度 - tags: - - SystemConfig - put: - consumes: - - application/json - description: 配置所有用户的并发任务上限、滚动窗口成功下载上限和豁免白名单 - parameters: - - description: 下载额度配置 - in: body - name: data - required: true - schema: - $ref: '#/definitions/internal_handler.UpdateModelDownloadLimitConfigReq' - produces: - - application/json - responses: - "200": - description: OK - schema: - $ref: '#/definitions/github_com_raids-lab_crater_internal_resputil.Response-string' - security: - - Bearer: [] - summary: 更新模型与数据集下载额度 - tags: - - SystemConfig /v1/admin/system-config/prequeue: get: description: 获取当前回填提交开关、Crater 队内资源配额开关、普通作业等待忍耐时间和 watcher 运行参数 @@ -6268,6 +6081,44 @@ paths: summary: 获取所有Cuda基础镜像 tags: - ImagePack + post: + consumes: + - application/json + description: 添加新的CUDA基础镜像到系统中 + parameters: + - description: CUDA基础镜像信息 + in: body + name: data + required: true + schema: + $ref: '#/definitions/internal_handler_image.CudaBaseImageCreateRequest' + produces: + - application/json + responses: {} + security: + - Bearer: [] + summary: 添加CUDA基础镜像 + tags: + - ImagePack + /v1/images/cudabaseimage/{id}: + delete: + consumes: + - application/json + description: 根据ID删除指定的CUDA基础镜像 + parameters: + - description: CUDA基础镜像ID + in: path + name: id + required: true + type: integer + produces: + - application/json + responses: {} + security: + - Bearer: [] + summary: 删除CUDA基础镜像 + tags: + - ImagePack /v1/images/deleteimage: post: consumes: @@ -6858,7 +6709,7 @@ paths: delete: consumes: - application/json - description: 删除下载任务记录(仅平台管理员),已下载的文件保留在存储中 + description: 删除下载任务记录(仅创建者或管理员),已下载的文件保留在存储中 parameters: - description: 下载任务ID in: path @@ -6984,7 +6835,7 @@ paths: name: id required: true type: integer - - description: 可选的临时访问令牌和重试版本 + - description: 可选的临时访问令牌 in: body name: data schema: @@ -8128,21 +7979,6 @@ paths: summary: 获取资源统计信息 tags: - statistics - /v1/system-config/model-download-limit: - get: - description: 获取当前用户的并发任务上限、滚动窗口成功下载上限和白名单豁免状态 - produces: - - application/json - responses: - "200": - description: OK - schema: - $ref: '#/definitions/github_com_raids-lab_crater_internal_resputil.Response-internal_handler_ModelDownloadLimitConfigResp' - security: - - Bearer: [] - summary: 获取模型与数据集下载额度 - tags: - - SystemConfig /v1/token/verify: get: consumes: diff --git a/backend/etc/example-config.yaml b/backend/etc/example-config.yaml index b5ee7dd27..d3663d46b 100644 --- a/backend/etc/example-config.yaml +++ b/backend/etc/example-config.yaml @@ -10,34 +10,10 @@ port: :8088 # Optional: If not specified, default values will be used modelDownload: # Container image used for model download jobs - # Optional: Defaults to Crater's public, version-pinned downloader image - # Deployments may mirror this image into an administrator-approved internal registry - # For a private registry, also set secrets.imagePullSecretName below - image: ghcr.io/raids-lab/crater-model-downloader:v1.0.0 - # Hugging Face Hub base URL used by download jobs - # Deployments may point this to an administrator-approved mirror or gateway - huggingFaceEndpoint: https://huggingface.co - # ModelScope base URL used by download jobs - modelScopeEndpoint: https://modelscope.cn - -# Source metadata refresh configuration -modelMetadata: - # Endpoints are tried in order; use only mirrors or gateways trusted by the administrator - huggingFaceEndpoints: - # - https://hf-mirror.com - - https://huggingface.co - modelScopeEndpoints: - - https://modelscope.cn - # Exact HTTPS hosts allowed when caching source logos; add trusted mirror/CDN hosts as needed - logoAllowedHosts: - - huggingface.co - - cdn-avatars.huggingface.co - - resouces.modelscope.cn - - resources.modelscope.cn - # Logical prefix in download records before mapping to storage.prefix.public - logicalPublicPrefix: public - timeoutSeconds: 20 - maxLogoBytes: 524288 + # Optional: Defaults to "python:3.11-slim" if not specified + # For deployments with network restrictions (e.g., in China), use internal registry: + # image: crater-harbor.act.buaa.edu.cn/crater/base/python:3.11-slim + image: python:3.11-slim # Endpoint URL for Prometheus API used for metrics and monitoring # Optional: If not specified, Prometheus integration will be disabled @@ -116,7 +92,6 @@ secrets: # Required: Secret must contain appropriate forwarding certificates tlsForwardSecretName: crater-tls-forward-secret # Name of the Kubernetes secret for pulling container images from private registries - # This is also applied to model download jobs when modelDownload.image is private # Optional: If not specified, no image pull secret will be used imagePullSecretName: "" @@ -155,10 +130,10 @@ registry: buildx: ghcr.io/raids-lab/buildx-client:latest # Nerdctl image for containerd-based builds # Required if Registry.Enable is true - nerdctl: ghcr.io/raids-lab/envd-client:latest + nerdctl: ghcr.io/raids-lab/nerdctl-client:latest # Envd image for environment-based development builds # Required if Registry.Enable is true - envd: ghcr.io/raids-lab/nerdctl-client:latest + envd: ghcr.io/raids-lab/envd-client:latest # Configuration for email notifications via SMTP # Optional: If Enable is false, email notifications will be disabled @@ -182,6 +157,12 @@ smtp: # Required if Enable is true: Must be a valid email address notify: example@example.com +# Crater Agent (Python service) +agent: + # Optional: defaults to http://localhost:8000 + serviceURL: http://localhost:8000 + # Optional: can be provided via CRATER_AGENT_INTERNAL_TOKEN + internalToken: # Authentication configuration auth: # Authentication token configuration for JWT-based authentication diff --git a/backend/go.mod b/backend/go.mod index 09d21c0c8..9bb8271a9 100644 --- a/backend/go.mod +++ b/backend/go.mod @@ -31,7 +31,6 @@ require ( gopkg.in/yaml.v3 v3.0.1 gorm.io/datatypes v1.2.6 gorm.io/driver/postgres v1.6.0 - gorm.io/driver/sqlite v1.6.0 gorm.io/gen v0.3.27 gorm.io/gorm v1.31.0 gorm.io/plugin/dbresolver v1.6.2 @@ -154,6 +153,7 @@ require ( gopkg.in/evanphx/json-patch.v4 v4.12.0 // indirect gopkg.in/inf.v0 v0.9.1 // indirect gorm.io/driver/mysql v1.5.7 // indirect + gorm.io/driver/sqlite v1.6.0 // indirect gorm.io/hints v1.1.0 // indirect k8s.io/apiextensions-apiserver v0.33.0 // indirect k8s.io/kube-openapi v0.0.0-20250318190949-c8a335a9a2ff // indirect diff --git a/backend/hack/sql/20260403_agent_multi_mode.sql b/backend/hack/sql/20260403_agent_multi_mode.sql new file mode 100644 index 000000000..16c11ff9f --- /dev/null +++ b/backend/hack/sql/20260403_agent_multi_mode.sql @@ -0,0 +1,128 @@ +BEGIN; + +ALTER TABLE agent_sessions + ADD COLUMN IF NOT EXISTS last_orchestration_mode VARCHAR(32) DEFAULT 'single_agent'; + +UPDATE agent_sessions +SET last_orchestration_mode = 'single_agent' +WHERE last_orchestration_mode IS NULL + OR BTRIM(last_orchestration_mode) = ''; + +ALTER TABLE agent_tool_calls + ADD COLUMN IF NOT EXISTS turn_id UUID, + ADD COLUMN IF NOT EXISTS tool_call_id VARCHAR(128), + ADD COLUMN IF NOT EXISTS agent_id VARCHAR(128), + ADD COLUMN IF NOT EXISTS parent_event_id BIGINT, + ADD COLUMN IF NOT EXISTS agent_role VARCHAR(32); + +CREATE INDEX IF NOT EXISTS idx_agent_tool_calls_turn_id ON agent_tool_calls(turn_id); +CREATE INDEX IF NOT EXISTS idx_agent_tool_calls_tool_call_id ON agent_tool_calls(tool_call_id); +CREATE INDEX IF NOT EXISTS idx_agent_tool_calls_agent_id ON agent_tool_calls(agent_id); +CREATE INDEX IF NOT EXISTS idx_agent_tool_calls_parent_event_id ON agent_tool_calls(parent_event_id); +CREATE INDEX IF NOT EXISTS idx_agent_tool_calls_agent_role ON agent_tool_calls(agent_role); + +CREATE TABLE IF NOT EXISTS agent_turns ( + id BIGSERIAL PRIMARY KEY, + turn_id UUID NOT NULL UNIQUE, + session_id UUID NOT NULL, + request_id VARCHAR(128), + orchestration_mode VARCHAR(32) NOT NULL DEFAULT 'single_agent', + root_agent_id VARCHAR(128), + status VARCHAR(32) NOT NULL DEFAULT 'running', + final_message_id BIGINT, + metadata JSONB, + started_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + ended_at TIMESTAMPTZ, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); + +ALTER TABLE agent_turns + ADD COLUMN IF NOT EXISTS turn_id UUID, + ADD COLUMN IF NOT EXISTS session_id UUID, + ADD COLUMN IF NOT EXISTS request_id VARCHAR(128), + ADD COLUMN IF NOT EXISTS orchestration_mode VARCHAR(32), + ADD COLUMN IF NOT EXISTS root_agent_id VARCHAR(128), + ADD COLUMN IF NOT EXISTS status VARCHAR(32), + ADD COLUMN IF NOT EXISTS final_message_id BIGINT, + ADD COLUMN IF NOT EXISTS metadata JSONB, + ADD COLUMN IF NOT EXISTS started_at TIMESTAMPTZ, + ADD COLUMN IF NOT EXISTS ended_at TIMESTAMPTZ, + ADD COLUMN IF NOT EXISTS created_at TIMESTAMPTZ, + ADD COLUMN IF NOT EXISTS updated_at TIMESTAMPTZ; + +ALTER TABLE agent_turns + ALTER COLUMN orchestration_mode SET DEFAULT 'single_agent', + ALTER COLUMN status SET DEFAULT 'running', + ALTER COLUMN started_at SET DEFAULT NOW(), + ALTER COLUMN created_at SET DEFAULT NOW(), + ALTER COLUMN updated_at SET DEFAULT NOW(); + +CREATE INDEX IF NOT EXISTS idx_agent_turns_session_id ON agent_turns(session_id); +CREATE UNIQUE INDEX IF NOT EXISTS idx_agent_turns_turn_id ON agent_turns(turn_id); +CREATE INDEX IF NOT EXISTS idx_agent_turns_request_id ON agent_turns(request_id); +CREATE INDEX IF NOT EXISTS idx_agent_turns_orchestration_mode ON agent_turns(orchestration_mode); +CREATE INDEX IF NOT EXISTS idx_agent_turns_root_agent_id ON agent_turns(root_agent_id); +CREATE INDEX IF NOT EXISTS idx_agent_turns_status ON agent_turns(status); +CREATE INDEX IF NOT EXISTS idx_agent_turns_final_message_id ON agent_turns(final_message_id); +CREATE INDEX IF NOT EXISTS idx_agent_turns_started_at ON agent_turns(started_at); + +CREATE TABLE IF NOT EXISTS agent_run_events ( + id BIGSERIAL PRIMARY KEY, + turn_id UUID NOT NULL, + session_id UUID NOT NULL, + agent_id VARCHAR(128), + parent_agent_id VARCHAR(128), + agent_role VARCHAR(32), + event_type VARCHAR(64) NOT NULL, + event_status VARCHAR(32), + title VARCHAR(255), + content TEXT, + metadata JSONB, + sequence INTEGER NOT NULL DEFAULT 0, + started_at TIMESTAMPTZ, + ended_at TIMESTAMPTZ, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); + +ALTER TABLE agent_run_events + ADD COLUMN IF NOT EXISTS turn_id UUID, + ADD COLUMN IF NOT EXISTS session_id UUID, + ADD COLUMN IF NOT EXISTS agent_id VARCHAR(128), + ADD COLUMN IF NOT EXISTS parent_agent_id VARCHAR(128), + ADD COLUMN IF NOT EXISTS agent_role VARCHAR(32), + ADD COLUMN IF NOT EXISTS event_type VARCHAR(64), + ADD COLUMN IF NOT EXISTS event_status VARCHAR(32), + ADD COLUMN IF NOT EXISTS title VARCHAR(255), + ADD COLUMN IF NOT EXISTS content TEXT, + ADD COLUMN IF NOT EXISTS metadata JSONB, + ADD COLUMN IF NOT EXISTS sequence INTEGER, + ADD COLUMN IF NOT EXISTS started_at TIMESTAMPTZ, + ADD COLUMN IF NOT EXISTS ended_at TIMESTAMPTZ, + ADD COLUMN IF NOT EXISTS created_at TIMESTAMPTZ; + +ALTER TABLE agent_run_events + ALTER COLUMN sequence SET DEFAULT 0, + ALTER COLUMN created_at SET DEFAULT NOW(); + +CREATE INDEX IF NOT EXISTS idx_agent_run_events_turn_id ON agent_run_events(turn_id); +CREATE INDEX IF NOT EXISTS idx_agent_run_events_session_id ON agent_run_events(session_id); +CREATE INDEX IF NOT EXISTS idx_agent_run_events_agent_id ON agent_run_events(agent_id); +CREATE INDEX IF NOT EXISTS idx_agent_run_events_parent_agent_id ON agent_run_events(parent_agent_id); +CREATE INDEX IF NOT EXISTS idx_agent_run_events_agent_role ON agent_run_events(agent_role); +CREATE INDEX IF NOT EXISTS idx_agent_run_events_event_type ON agent_run_events(event_type); +CREATE INDEX IF NOT EXISTS idx_agent_run_events_event_status ON agent_run_events(event_status); +CREATE INDEX IF NOT EXISTS idx_agent_run_events_sequence ON agent_run_events(sequence); +CREATE INDEX IF NOT EXISTS idx_agent_run_events_turn_sequence ON agent_run_events(turn_id, sequence); + +COMMIT; + +-- Add pinned_at column for session pinning feature +BEGIN; + +ALTER TABLE agent_sessions + ADD COLUMN IF NOT EXISTS pinned_at TIMESTAMPTZ; + +CREATE INDEX IF NOT EXISTS idx_agent_sessions_pinned_at ON agent_sessions(pinned_at); + +COMMIT; diff --git a/backend/hack/sql/20260422_agent_session_source.sql b/backend/hack/sql/20260422_agent_session_source.sql new file mode 100644 index 000000000..f9780367b --- /dev/null +++ b/backend/hack/sql/20260422_agent_session_source.sql @@ -0,0 +1,22 @@ +-- Add source field to agent_sessions to distinguish chat vs non-chat sessions. +-- Values: 'chat' (default, user-initiated conversations shown in UI), +-- 'ops_audit' (admin tool execution audit, hidden from chat UI), +-- 'system' (background/automated agent operations). +ALTER TABLE agent_sessions ADD COLUMN IF NOT EXISTS source VARCHAR(32) NOT NULL DEFAULT 'chat'; +UPDATE agent_sessions +SET source = 'chat' +WHERE source IS NULL OR BTRIM(source) = ''; +UPDATE agent_sessions +SET source = 'ops_audit' +WHERE title LIKE '[audit] 审批%' + AND source = 'chat'; +CREATE INDEX IF NOT EXISTS idx_agent_sessions_source ON agent_sessions (source); + +-- Add source field to agent_tool_calls to track execution origin. +-- Values: 'backend' (default, executed via Go backend API), +-- 'local' (executed locally by Python agent via kubectl/prometheus). +ALTER TABLE agent_tool_calls ADD COLUMN IF NOT EXISTS source VARCHAR(32) NOT NULL DEFAULT 'backend'; +UPDATE agent_tool_calls +SET source = 'backend' +WHERE source IS NULL OR BTRIM(source) = ''; +CREATE INDEX IF NOT EXISTS idx_agent_tool_calls_source ON agent_tool_calls (source); diff --git a/backend/internal/governance/modeldataset/provenance.go b/backend/internal/governance/modeldataset/provenance.go deleted file mode 100644 index bdf179f15..000000000 --- a/backend/internal/governance/modeldataset/provenance.go +++ /dev/null @@ -1,251 +0,0 @@ -// Copyright 2026 The Crater Project Team, RAIDS-Lab -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package modeldataset - -import ( - "encoding/json" - "errors" - "io" - "os" - "path/filepath" - "reflect" - "regexp" - "strconv" - "strings" - - "github.com/raids-lab/crater/dao/model" -) - -const maxProvenanceFileBytes = 256 * 1024 - -const ( - provenanceConfidenceHigh = "high" - provenanceConfidenceMedium = "medium" -) - -var ( - huggingFaceURLPattern = regexp.MustCompile(`https?://(?:www\.)?(?:huggingface\.co|hf-mirror\.com)/([A-Za-z0-9_.-]+/[A-Za-z0-9_.-]+)`) - modelScopeURLPattern = regexp.MustCompile(`https?://(?:www\.)?modelscope\.cn/(?:models|datasets)/([A-Za-z0-9_.-]+/[A-Za-z0-9_.-]+)`) - huggingFaceSSHPattern = regexp.MustCompile(`(?:git@|ssh://git@)hf\.co[:/]([A-Za-z0-9_.-]+/[A-Za-z0-9_.-]+)`) -) - -type repositoryReference struct { - Provider model.ModelDatasetProvider - RepositoryID string - URL string -} - -func collectFilesystemEvidence(directory string, evidence *model.ModelDatasetDiscoveryEvidence) error { - info, err := os.Stat(directory) - if err != nil { - return err - } - modifiedAt := info.ModTime() - evidence.ModifiedAt = &modifiedAt - evidence.FilesystemUID, evidence.FilesystemGID = filesystemOwnership(info) - - configNameOrPath, err := readConfigNameOrPath(filepath.Join(directory, "config.json")) - if err != nil { - return err - } - evidence.ConfigNameOrPath = configNameOrPath - - gitReferences, err := referencesFromFile(filepath.Join(directory, ".git", "config")) - if err != nil { - return err - } - if reference, ok := oneReference(gitReferences); ok { - applyRepositoryEvidence(evidence, reference, "git_remote", provenanceConfidenceHigh) - return nil - } - - readmeReferences := make([]repositoryReference, 0) - for _, name := range []string{"README.md", "readme.md", "README.MD", "README"} { - references, readErr := referencesFromFile(filepath.Join(directory, name)) - if readErr != nil { - return readErr - } - readmeReferences = append(readmeReferences, references...) - } - readmeReferences = uniqueReferences(readmeReferences) - for _, reference := range readmeReferences { - evidence.CandidateURLs = append(evidence.CandidateURLs, reference.URL) - } - if reference, ok := matchingReadmeReference(readmeReferences, filepath.Base(directory), configNameOrPath); ok { - applyRepositoryEvidence(evidence, reference, "readme_url", provenanceConfidenceHigh) - return nil - } - - if validRepositoryID(configNameOrPath) { - evidence.RepositoryID = strings.TrimSuffix(configNameOrPath, ".git") - evidence.ProvenanceSource = "config_name_or_path" - evidence.ProvenanceConfidence = provenanceConfidenceMedium - } - return nil -} - -// filesystemOwnership uses the platform-provided stat object without tying the -// scanner to one operating system. UID/GID remain optional evidence on systems -// that do not expose Unix ownership. -func filesystemOwnership(info os.FileInfo) (uid, gid string) { - stat := reflect.ValueOf(info.Sys()) - if !stat.IsValid() { - return "", "" - } - if stat.Kind() == reflect.Pointer { - stat = stat.Elem() - } - if !stat.IsValid() || stat.Kind() != reflect.Struct { - return "", "" - } - value := func(name string) string { - field := stat.FieldByName(name) - if !field.IsValid() { - return "" - } - switch field.Kind() { - case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64: - return strconv.FormatUint(field.Uint(), 10) - case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64: - return strconv.FormatInt(field.Int(), 10) - default: - return "" - } - } - return value("Uid"), value("Gid") -} - -func readConfigNameOrPath(path string) (string, error) { - data, err := readSmallFile(path) - if err != nil { - return "", err - } - if len(data) == 0 { - return "", nil - } - var config struct { - NameOrPath string `json:"_name_or_path"` - } - if err := json.Unmarshal(data, &config); err != nil { - // A malformed config must not prevent inventorying an otherwise complete model. - return "", nil - } - return strings.TrimSpace(config.NameOrPath), nil -} - -func referencesFromFile(path string) ([]repositoryReference, error) { - data, err := readSmallFile(path) - if err != nil { - return nil, err - } - if len(data) == 0 { - return nil, nil - } - return referencesFromText(string(data)), nil -} - -func readSmallFile(path string) ([]byte, error) { - file, err := os.Open(path) - if errors.Is(err, os.ErrNotExist) { - return nil, nil - } - if err != nil { - return nil, err - } - defer file.Close() - data, err := io.ReadAll(io.LimitReader(file, maxProvenanceFileBytes)) - if err != nil { - return nil, err - } - return data, nil -} - -func referencesFromText(text string) []repositoryReference { - result := make([]repositoryReference, 0) - appendMatches := func(pattern *regexp.Regexp, provider model.ModelDatasetProvider, canonicalBase string) { - for _, match := range pattern.FindAllStringSubmatch(text, -1) { - repositoryID := strings.TrimSuffix(strings.TrimRight(match[1], ").,;]}'\""), ".git") - result = append(result, repositoryReference{ - Provider: provider, RepositoryID: repositoryID, URL: canonicalBase + repositoryID, - }) - } - } - appendMatches(huggingFaceURLPattern, model.ModelDatasetProviderHuggingFace, "https://huggingface.co/") - appendMatches(huggingFaceSSHPattern, model.ModelDatasetProviderHuggingFace, "https://huggingface.co/") - appendMatches(modelScopeURLPattern, model.ModelDatasetProviderModelScope, "https://modelscope.cn/models/") - return uniqueReferences(result) -} - -func uniqueReferences(references []repositoryReference) []repositoryReference { - result := make([]repositoryReference, 0, len(references)) - seen := make(map[string]struct{}) - for _, reference := range references { - key := string(reference.Provider) + "\x00" + reference.RepositoryID - if _, exists := seen[key]; exists { - continue - } - seen[key] = struct{}{} - result = append(result, reference) - } - return result -} - -func oneReference(references []repositoryReference) (repositoryReference, bool) { - references = uniqueReferences(references) - returnReference := repositoryReference{} - if len(references) == 1 { - returnReference = references[0] - return returnReference, true - } - return returnReference, false -} - -func matchingReadmeReference( - references []repositoryReference, - directoryName, configNameOrPath string, -) (repositoryReference, bool) { - matches := make([]repositoryReference, 0) - for _, reference := range references { - parts := strings.Split(reference.RepositoryID, "/") - baseMatches := len(parts) == 2 && strings.EqualFold(parts[1], directoryName) - configMatches := validRepositoryID(configNameOrPath) && - strings.EqualFold(strings.TrimSuffix(configNameOrPath, ".git"), reference.RepositoryID) - if baseMatches || configMatches { - matches = append(matches, reference) - } - } - return oneReference(matches) -} - -func validRepositoryID(value string) bool { - value = strings.TrimSpace(strings.TrimSuffix(value, ".git")) - if value == "" || filepath.IsAbs(value) || strings.Contains(value, "\\") || strings.Contains(value, ":") { - return false - } - parts := strings.Split(value, "/") - return len(parts) == 2 && parts[0] != "" && parts[1] != "" && parts[0] != "." && parts[0] != ".." -} - -func applyRepositoryEvidence( - evidence *model.ModelDatasetDiscoveryEvidence, - reference repositoryReference, - source, confidence string, -) { - evidence.Provider = reference.Provider - evidence.RepositoryID = reference.RepositoryID - evidence.RepositoryURL = reference.URL - evidence.ProvenanceSource = source - evidence.ProvenanceConfidence = confidence -} diff --git a/backend/internal/governance/modeldataset/reconcile.go b/backend/internal/governance/modeldataset/reconcile.go deleted file mode 100644 index ff91dbaca..000000000 --- a/backend/internal/governance/modeldataset/reconcile.go +++ /dev/null @@ -1,641 +0,0 @@ -// Copyright 2026 The Crater Project Team, RAIDS-Lab -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package modeldataset - -import ( - "context" - "errors" - "fmt" - "os" - "path/filepath" - "strconv" - "strings" - "time" - "unicode/utf8" - - "gorm.io/datatypes" - "gorm.io/gorm" - - "github.com/raids-lab/crater/dao/model" -) - -type ReconcileOptions struct { - Apply bool - LogicalPublicPrefix string - PhysicalPublicPrefix string - PhysicalUserPrefix string - PhysicalAccountPrefix string - MaxReadmeBytes int - Now time.Time -} - -type ReconcileReport struct { - ReadyDownloads int - PublicDownloads int - SourceIdentities int - DownloadLinks int - DatasetLinks int - Candidates int - RegisteredCandidates int - UnregisteredCandidates int - ReadmesFromStorage int - DiscoveriesWritten int - RegisteredNonPublic int - PathlessDatasets int - MissingMarked int64 - ProvenanceHighConfidence int - ProvenanceHints int - InferredSources int - OwnerMatches int -} - -type sourceIdentity struct { - Provider model.ModelDatasetProvider - ResourceType model.DataType - RepositoryID string -} - -//nolint:gocyclo,funlen // Reconciliation reports and applies one idempotent plan in a single pass. -func ReconcilePublic( - ctx context.Context, - db *gorm.DB, - candidates []Candidate, - options *ReconcileOptions, -) (ReconcileReport, error) { - if options.LogicalPublicPrefix == "" || options.PhysicalPublicPrefix == "" { - return ReconcileReport{}, errors.New("logical and physical public prefixes are required") - } - if options.MaxReadmeBytes < 1 { - return ReconcileReport{}, errors.New("max README bytes must be positive") - } - if options.Now.IsZero() { - options.Now = time.Now() - } - - var downloads []*model.ModelDownload - if err := db.WithContext(ctx). - Where("status = ?", model.ModelDownloadStatusReady). - Order("id ASC"). - Find(&downloads).Error; err != nil { - return ReconcileReport{}, fmt.Errorf("load ready downloads: %w", err) - } - var datasets []*model.Dataset - if err := db.WithContext(ctx). - Where("deleted_at IS NULL AND type IN ?", []model.DataType{model.DataTypeModel, model.DataTypeDataset}). - Find(&datasets).Error; err != nil { - return ReconcileReport{}, fmt.Errorf("load datasets: %w", err) - } - - report := ReconcileReport{ReadyDownloads: len(downloads), Candidates: len(candidates)} - datasetsByPath := make(map[string][]*model.Dataset) - datasetsByResource := make(map[string][]*model.Dataset) - for _, dataset := range datasets { - datasetsByPath[cleanStoragePath(dataset.URL)] = append(datasetsByPath[cleanStoragePath(dataset.URL)], dataset) - key := resourceKey(dataset.Name, dataset.Type) - datasetsByResource[key] = append(datasetsByResource[key], dataset) - } - ownerMatches, err := enrichCandidateOwners(ctx, db, candidates) - if err != nil { - return report, err - } - report.OwnerMatches = ownerMatches - candidatesByPath := make(map[string]*Candidate, len(candidates)) - for index := range candidates { - candidate := &candidates[index] - candidatesByPath[cleanStoragePath(candidate.Path)] = candidate - } - - seenIdentities := make(map[sourceIdentity]struct{}) - sourceByPath := make(map[string]*model.ModelDatasetSource) - sourceIDs := make([]uint, 0) - for _, dataset := range datasets { - if dataset.ModelDatasetSourceID != nil { - sourceIDs = append(sourceIDs, *dataset.ModelDatasetSourceID) - } - } - if len(sourceIDs) > 0 { - var persistedSources []*model.ModelDatasetSource - if err := db.WithContext(ctx).Where("id IN ?", sourceIDs).Find(&persistedSources).Error; err != nil { - return report, fmt.Errorf("load linked sources: %w", err) - } - sourcesByID := make(map[uint]*model.ModelDatasetSource, len(persistedSources)) - for _, source := range persistedSources { - sourcesByID[source.ID] = source - } - for _, dataset := range datasets { - if dataset.ModelDatasetSourceID == nil { - continue - } - if source := sourcesByID[*dataset.ModelDatasetSourceID]; source != nil { - sourceByPath[cleanStoragePath(dataset.URL)] = source - seenIdentities[sourceIdentity{ - Provider: source.Provider, ResourceType: source.ResourceType, RepositoryID: source.RepositoryID, - }] = struct{}{} - } - } - } - for _, download := range downloads { - physicalPath, public := PhysicalStoragePath( - download.Path, - options.LogicalPublicPrefix, - options.PhysicalPublicPrefix, - ) - if !public { - continue - } - report.PublicDownloads++ - identity := sourceIdentity{ - Provider: model.ModelDatasetProvider(download.Source), - ResourceType: model.DataType(download.Category), - RepositoryID: download.Name, - } - seenIdentities[identity] = struct{}{} - - source := sourceFromDownload(download) - if candidate, ok := candidatesByPath[physicalPath]; ok { - readme, err := readLocalReadme(candidate.AbsolutePath, options.MaxReadmeBytes) - if err != nil { - return report, fmt.Errorf("read README for %s: %w", physicalPath, err) - } - if readme != "" { - source.Readme = readme - report.ReadmesFromStorage++ - } - } - - if options.Apply { - persisted, err := upsertSource(ctx, db, source) - if err != nil { - return report, err - } - source = persisted - if err := db.WithContext(ctx).Model(&model.ModelDownload{}). - Where("id = ?", download.ID). - Update("model_dataset_source_id", source.ID).Error; err != nil { - return report, fmt.Errorf("link download %d to source: %w", download.ID, err) - } - report.DownloadLinks++ - linkedDataset := false - for _, dataset := range datasetsByPath[physicalPath] { - if dataset.Type != identity.ResourceType { - continue - } - if err := db.WithContext(ctx).Model(&model.Dataset{}). - Where("id = ?", dataset.ID). - Update("model_dataset_source_id", source.ID).Error; err != nil { - return report, fmt.Errorf("link dataset %d to source: %w", dataset.ID, err) - } - dataset.ModelDatasetSourceID = &source.ID - report.DatasetLinks++ - linkedDataset = true - } - // Historical shared records may have an empty or non-standard path. Link only - // when repository identity and type select one unambiguous Dataset row. - if !linkedDataset { - matches := datasetsByResource[resourceKey(download.Name, identity.ResourceType)] - if len(matches) == 1 && matches[0].ModelDatasetSourceID == nil { - dataset := matches[0] - if err := db.WithContext(ctx).Model(&model.Dataset{}). - Where("id = ?", dataset.ID). - Update("model_dataset_source_id", source.ID).Error; err != nil { - return report, fmt.Errorf("link pathless dataset %d to source: %w", dataset.ID, err) - } - dataset.ModelDatasetSourceID = &source.ID - report.DatasetLinks++ - } - } - } - sourceByPath[physicalPath] = source - } - seenPaths := make([]string, 0, len(candidates)) - for index := range candidates { - candidate := &candidates[index] - path := cleanStoragePath(candidate.Path) - seenPaths = append(seenPaths, path) - if candidate.Evidence.ProvenanceConfidence == provenanceConfidenceHigh { - report.ProvenanceHighConfidence++ - } else if candidate.Evidence.RepositoryID != "" || candidate.Evidence.ConfigNameOrPath != "" || - len(candidate.Evidence.CandidateURLs) > 0 { - report.ProvenanceHints++ - } - status := model.ModelDatasetDiscoveryStatusDiscovered - var datasetID, sourceID *uint - if matches := datasetsByPath[path]; len(matches) > 0 { - status = model.ModelDatasetDiscoveryStatusRegistered - datasetID = &matches[0].ID - report.RegisteredCandidates++ - } else { - report.UnregisteredCandidates++ - } - source := sourceByPath[path] - if source == nil && canInferSource(&candidate.Evidence) { - report.InferredSources++ - identity := sourceIdentity{ - Provider: candidate.Evidence.Provider, ResourceType: candidate.Type, - RepositoryID: candidate.Evidence.RepositoryID, - } - seenIdentities[identity] = struct{}{} - inferred := sourceFromCandidate(candidate) - if options.Apply { - readme, err := readLocalReadme(candidate.AbsolutePath, options.MaxReadmeBytes) - if err != nil { - return report, fmt.Errorf("read README for %s: %w", path, err) - } - if readme != "" { - inferred.Readme = readme - report.ReadmesFromStorage++ - } - persisted, err := upsertSource(ctx, db, inferred) - if err != nil { - return report, err - } - inferred = persisted - for _, dataset := range datasetsByPath[path] { - if dataset.Type != candidate.Type || dataset.ModelDatasetSourceID != nil { - continue - } - if err := db.WithContext(ctx).Model(&model.Dataset{}).Where("id = ?", dataset.ID). - Update("model_dataset_source_id", inferred.ID).Error; err != nil { - return report, fmt.Errorf("link dataset %d to inferred source: %w", dataset.ID, err) - } - dataset.ModelDatasetSourceID = &inferred.ID - report.DatasetLinks++ - } - } - source = inferred - sourceByPath[path] = inferred - } - if source != nil && source.ID != 0 { - sourceID = &source.ID - } - if !options.Apply { - continue - } - discovery := model.ModelDatasetDiscovery{ - DiscoveryKey: "path:" + path, - Path: path, - Scope: model.ModelDatasetDiscoveryScopePublic, - DetectedType: candidate.Type, - DetectedName: candidate.Name, - Evidence: datatypes.NewJSONType(candidate.Evidence), - SizeBytes: candidate.SizeBytes, - DatasetID: datasetID, - SourceID: sourceID, - Status: status, - FirstSeenAt: options.Now, - LastSeenAt: options.Now, - } - if err := upsertDiscovery(ctx, db, &discovery); err != nil { - return report, err - } - report.DiscoveriesWritten++ - } - - for _, dataset := range datasets { - path := cleanStoragePath(dataset.URL) - if path == "" { - report.PathlessDatasets++ - if options.Apply { - datasetID := dataset.ID - ownerID := dataset.UserID - discovery := model.ModelDatasetDiscovery{ - DiscoveryKey: "dataset:" + strconv.FormatUint(uint64(dataset.ID), 10), - Scope: model.ModelDatasetDiscoveryScopeUser, ScopeID: &ownerID, - DetectedType: dataset.Type, DetectedName: dataset.Name, - DatasetID: &datasetID, SourceID: dataset.ModelDatasetSourceID, - Status: model.ModelDatasetDiscoveryStatusPathMissing, - FirstSeenAt: options.Now, LastSeenAt: options.Now, - } - if err := upsertDiscovery(ctx, db, &discovery); err != nil { - return report, err - } - report.DiscoveriesWritten++ - } - continue - } - if pathUnderPrefix(path, options.PhysicalPublicPrefix) { - continue - } - scope := model.ModelDatasetDiscoveryScopeUser - var scopeID *uint - if pathUnderPrefix(path, options.PhysicalAccountPrefix) { - scope = model.ModelDatasetDiscoveryScopeAccount - } else { - ownerID := dataset.UserID - scopeID = &ownerID - } - report.RegisteredNonPublic++ - if !options.Apply { - continue - } - datasetID := dataset.ID - discovery := model.ModelDatasetDiscovery{ - DiscoveryKey: "dataset:" + strconv.FormatUint(uint64(dataset.ID), 10), - Path: path, Scope: scope, ScopeID: scopeID, - DetectedType: dataset.Type, DetectedName: dataset.Name, - DatasetID: &datasetID, SourceID: dataset.ModelDatasetSourceID, - Status: model.ModelDatasetDiscoveryStatusRegistered, - FirstSeenAt: options.Now, LastSeenAt: options.Now, - } - if err := upsertDiscovery(ctx, db, &discovery); err != nil { - return report, err - } - report.DiscoveriesWritten++ - } - - if options.Apply { - query := db.WithContext(ctx).Model(&model.ModelDatasetDiscovery{}). - Where("scope = ?", model.ModelDatasetDiscoveryScopePublic). - Where("path = ? OR path LIKE ?", options.PhysicalPublicPrefix, options.PhysicalPublicPrefix+"/%") - if len(seenPaths) > 0 { - query = query.Where("path NOT IN ?", seenPaths) - } - result := query.Updates(map[string]any{ - "status": model.ModelDatasetDiscoveryStatusMissing, - "last_seen_at": options.Now, - }) - if result.Error != nil { - return report, fmt.Errorf("mark missing discoveries: %w", result.Error) - } - report.MissingMarked = result.RowsAffected - } - report.SourceIdentities = len(seenIdentities) - return report, nil -} - -func canInferSource(evidence *model.ModelDatasetDiscoveryEvidence) bool { - return evidence.ProvenanceConfidence == provenanceConfidenceHigh && evidence.RepositoryID != "" && - (evidence.Provider == model.ModelDatasetProviderHuggingFace || - evidence.Provider == model.ModelDatasetProviderModelScope) -} - -func sourceFromCandidate(candidate *Candidate) *model.ModelDatasetSource { - organization := "" - if parts := strings.Split(candidate.Evidence.RepositoryID, "/"); len(parts) == 2 { - organization = parts[0] - } - return &model.ModelDatasetSource{ - Provider: candidate.Evidence.Provider, ResourceType: candidate.Type, - RepositoryID: candidate.Evidence.RepositoryID, RepositoryURL: candidate.Evidence.RepositoryURL, - Organization: organization, DisplayName: candidate.Name, - } -} - -func enrichCandidateOwners(ctx context.Context, db *gorm.DB, candidates []Candidate) (int, error) { - if !db.Migrator().HasTable(&model.User{}) { - return 0, nil - } - var users []*model.User - if err := db.WithContext(ctx).Select("id", "name", "attributes").Find(&users).Error; err != nil { - return 0, fmt.Errorf("load filesystem user identities: %w", err) - } - usersByUID := make(map[string][]*model.User) - for _, user := range users { - attributes := user.Attributes.Data() - if attributes.UID == nil { - continue - } - uid := strings.TrimSpace(*attributes.UID) - if uid != "" { - usersByUID[uid] = append(usersByUID[uid], user) - } - } - matches := 0 - for index := range candidates { - candidate := &candidates[index] - users := usersByUID[candidate.Evidence.FilesystemUID] - if len(users) != 1 { - continue - } - ownerID := users[0].ID - candidate.Evidence.OwnerUserID = &ownerID - candidate.Evidence.OwnerUsername = users[0].Name - matches++ - } - return matches, nil -} - -func resourceKey(name string, resourceType model.DataType) string { - return string(resourceType) + "\x00" + name -} - -func pathUnderPrefix(path, prefix string) bool { - prefix = cleanStoragePath(prefix) - if prefix == "" { - return false - } - return path == prefix || strings.HasPrefix(path, prefix+"/") -} - -func PhysicalStoragePath(path, logicalPublicPrefix, physicalPublicPrefix string) (string, bool) { - path = cleanStoragePath(path) - logicalPublicPrefix = cleanStoragePath(logicalPublicPrefix) - physicalPublicPrefix = cleanStoragePath(physicalPublicPrefix) - if path == logicalPublicPrefix { - return physicalPublicPrefix, true - } - if strings.HasPrefix(path, logicalPublicPrefix+"/") { - return physicalPublicPrefix + strings.TrimPrefix(path, logicalPublicPrefix), true - } - if path == physicalPublicPrefix || strings.HasPrefix(path, physicalPublicPrefix+"/") { - return path, true - } - return path, false -} - -// LogicalStoragePath converts a physical storage prefix back to the stable logical -// prefix used by the file-system API and user-facing mount paths. -func LogicalStoragePath(path, logicalPrefix, physicalPrefix string) (string, bool) { - path = cleanStoragePath(path) - logicalPrefix = cleanStoragePath(logicalPrefix) - physicalPrefix = cleanStoragePath(physicalPrefix) - if path == logicalPrefix || strings.HasPrefix(path, logicalPrefix+"/") { - return path, true - } - if path == physicalPrefix { - return logicalPrefix, true - } - if strings.HasPrefix(path, physicalPrefix+"/") { - return logicalPrefix + strings.TrimPrefix(path, physicalPrefix), true - } - return path, false -} - -func sourceFromDownload(download *model.ModelDownload) *model.ModelDatasetSource { - return &model.ModelDatasetSource{ - Provider: model.ModelDatasetProvider(download.Source), - ResourceType: model.DataType(download.Category), - RepositoryID: download.Name, - RepositoryURL: download.SourceURL, - Organization: download.Organization, - LogoURL: download.LogoURL, - DisplayName: download.DisplayName, - Description: download.SourceDescription, - Readme: download.SourceReadme, - License: download.License, - Task: download.Task, - Library: download.Library, - ModelType: download.ModelType, - ParameterCount: download.ParameterCount, - Private: download.SourcePrivate, - Gated: download.SourceGated, - LoginRequired: download.SourceLoginRequired, - Downloads: download.SourceDownloads, - Likes: download.SourceLikes, - SourceCreatedAt: download.SourceCreatedAt, - SourceUpdatedAt: download.SourceUpdatedAt, - MetadataRefreshedAt: download.MetadataRefreshedAt, - } -} - -func upsertSource( - ctx context.Context, - db *gorm.DB, - source *model.ModelDatasetSource, -) (*model.ModelDatasetSource, error) { - var existing model.ModelDatasetSource - err := db.WithContext(ctx).Where( - "provider = ? AND resource_type = ? AND repository_id = ?", - source.Provider, source.ResourceType, source.RepositoryID, - ).First(&existing).Error - if errors.Is(err, gorm.ErrRecordNotFound) { - if err := db.WithContext(ctx).Create(source).Error; err != nil { - return nil, fmt.Errorf("create source %s/%s/%s: %w", source.Provider, source.ResourceType, source.RepositoryID, err) - } - return source, nil - } - if err != nil { - return nil, fmt.Errorf("load source %s/%s/%s: %w", source.Provider, source.ResourceType, source.RepositoryID, err) - } - updates := nonEmptySourceUpdates(source) - if len(updates) > 0 { - if err := db.WithContext(ctx).Model(&existing).Updates(updates).Error; err != nil { - return nil, fmt.Errorf("update source %d: %w", existing.ID, err) - } - } - return &existing, nil -} - -func nonEmptySourceUpdates(source *model.ModelDatasetSource) map[string]any { - updates := make(map[string]any) - stringsByColumn := map[string]string{ - "repository_url": source.RepositoryURL, - "organization": source.Organization, - "logo_url": source.LogoURL, - "display_name": source.DisplayName, - "description": source.Description, - "readme": source.Readme, - "license": source.License, - "task": source.Task, - "library": source.Library, - "model_type": source.ModelType, - } - for column, value := range stringsByColumn { - if value != "" { - updates[column] = value - } - } - if source.ParameterCount > 0 { - updates["parameter_count"] = source.ParameterCount - } - if source.Downloads > 0 { - updates["downloads"] = source.Downloads - } - if source.Likes > 0 { - updates["likes"] = source.Likes - } - if source.Private { - updates["private"] = true - } - if source.Gated { - updates["gated"] = true - } - if source.LoginRequired { - updates["login_required"] = true - } - if source.SourceCreatedAt != nil { - updates["source_created_at"] = source.SourceCreatedAt - } - if source.SourceUpdatedAt != nil { - updates["source_updated_at"] = source.SourceUpdatedAt - } - if source.MetadataRefreshedAt != nil { - updates["metadata_refreshed_at"] = source.MetadataRefreshedAt - } - return updates -} - -func upsertDiscovery(ctx context.Context, db *gorm.DB, discovery *model.ModelDatasetDiscovery) error { - var existing model.ModelDatasetDiscovery - err := db.WithContext(ctx).Where("discovery_key = ?", discovery.DiscoveryKey).First(&existing).Error - if errors.Is(err, gorm.ErrRecordNotFound) { - if err := db.WithContext(ctx).Create(discovery).Error; err != nil { - return fmt.Errorf("create discovery %s: %w", discovery.Path, err) - } - return nil - } - if err != nil { - return fmt.Errorf("load discovery %s: %w", discovery.Path, err) - } - if err := db.WithContext(ctx).Model(&existing).Updates(map[string]any{ - "path": discovery.Path, - "scope": discovery.Scope, - "scope_id": discovery.ScopeID, - "detected_type": discovery.DetectedType, - "detected_name": discovery.DetectedName, - "evidence": discovery.Evidence, - "size_bytes": discovery.SizeBytes, - "dataset_id": discovery.DatasetID, - "source_id": discovery.SourceID, - "status": discovery.Status, - "last_seen_at": discovery.LastSeenAt, - }).Error; err != nil { - return fmt.Errorf("update discovery %s: %w", discovery.Path, err) - } - return nil -} - -func readLocalReadme(directory string, maxBytes int) (string, error) { - for _, name := range []string{"README.md", "readme.md", "README.MD", "README"} { - path := filepath.Join(directory, name) - data, err := os.ReadFile(path) - if errors.Is(err, os.ErrNotExist) { - continue - } - if err != nil { - return "", err - } - if len(data) > maxBytes { - data = data[:maxBytes] - for len(data) > 0 && !utf8.Valid(data) { - data = data[:len(data)-1] - } - } - return string(data), nil - } - return "", nil -} - -func cleanStoragePath(path string) string { - path = strings.TrimSpace(path) - if path == "" { - return "" - } - path = strings.TrimPrefix(filepath.ToSlash(filepath.Clean(filepath.FromSlash(path))), "/") - if path == "." { - return "" - } - return path -} diff --git a/backend/internal/governance/modeldataset/reconcile_test.go b/backend/internal/governance/modeldataset/reconcile_test.go deleted file mode 100644 index 543afdb0b..000000000 --- a/backend/internal/governance/modeldataset/reconcile_test.go +++ /dev/null @@ -1,284 +0,0 @@ -// Copyright 2026 The Crater Project Team, RAIDS-Lab -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package modeldataset - -import ( - "context" - "fmt" - "path/filepath" - "testing" - "time" - - "gorm.io/datatypes" - "gorm.io/driver/sqlite" - "gorm.io/gorm" - - "github.com/raids-lab/crater/dao/model" -) - -//nolint:gocyclo // The integration assertion covers links, local README, and personal inventory together. -func TestReconcilePublicLinksOnlyExactPhysicalPaths(t *testing.T) { - db, err := gorm.Open(sqlite.Open("file::memory:?cache=shared"), &gorm.Config{}) - if err != nil { - t.Fatal(err) - } - if err := db.AutoMigrate( - &model.ModelDatasetSource{}, - &model.Dataset{}, - &model.ModelDownload{}, - &model.ModelDatasetDiscovery{}, - ); err != nil { - t.Fatal(err) - } - download := model.ModelDownload{ - Name: "owner/model", - Source: model.ModelSourceHuggingFace, - Category: model.DownloadCategoryModel, - Path: "public/Models/owner/model/huggingface/default", - Status: model.ModelDownloadStatusReady, - CreatorID: 1, - } - if err := db.Create(&download).Error; err != nil { - t.Fatal(err) - } - exact := model.Dataset{ - Name: "owner/model", URL: "shared/Models/owner/model/huggingface/default", Type: model.DataTypeModel, - } - mismatch := model.Dataset{ - Name: "owner/model", URL: "homes/user/model", Type: model.DataTypeModel, - } - if err := db.Create(&exact).Error; err != nil { - t.Fatal(err) - } - if err := db.Create(&mismatch).Error; err != nil { - t.Fatal(err) - } - - directory := t.TempDir() - mustWriteFile(t, filepath.Join(directory, "README.md"), "# local README") - now := time.Now().Round(time.Second) - report, err := ReconcilePublic(context.Background(), db, []Candidate{{ - Path: exact.URL, - AbsolutePath: directory, - Type: model.DataTypeModel, - Name: "model", - Evidence: model.ModelDatasetDiscoveryEvidence{HasConfig: true, HasReadme: true, WeightFiles: 1}, - }}, &ReconcileOptions{ - Apply: true, - LogicalPublicPrefix: "public", - PhysicalPublicPrefix: "shared", - MaxReadmeBytes: 1024, - Now: now, - }) - if err != nil { - t.Fatalf("ReconcilePublic() error = %v", err) - } - if report.DatasetLinks != 1 || report.ReadmesFromStorage != 1 { - t.Fatalf("report = %#v", report) - } - if err := db.First(&exact, exact.ID).Error; err != nil { - t.Fatal(err) - } - if err := db.First(&mismatch, mismatch.ID).Error; err != nil { - t.Fatal(err) - } - if exact.ModelDatasetSourceID == nil { - t.Fatal("exact path dataset was not linked") - } - if mismatch.ModelDatasetSourceID != nil { - t.Fatal("same-name dataset at another path must not be linked") - } - var personalDiscovery model.ModelDatasetDiscovery - if err := db.Where("path = ?", mismatch.URL).First(&personalDiscovery).Error; err != nil { - t.Fatalf("personal dataset path was not inventoried: %v", err) - } - if personalDiscovery.Scope != model.ModelDatasetDiscoveryScopeUser || - personalDiscovery.DatasetID == nil || *personalDiscovery.DatasetID != mismatch.ID { - t.Fatalf("personal discovery = %#v", personalDiscovery) - } - var source model.ModelDatasetSource - if err := db.First(&source, *exact.ModelDatasetSourceID).Error; err != nil { - t.Fatal(err) - } - if source.Readme != "# local README" { - t.Fatalf("source.Readme = %q", source.Readme) - } -} - -func TestPhysicalStoragePathIsConfigurable(t *testing.T) { - physical, public := PhysicalStoragePath("catalog/Models/a/b", "catalog", "shared-assets") - if !public || physical != "shared-assets/Models/a/b" { - t.Fatalf("PhysicalStoragePath() = %q, %v", physical, public) - } - physical, public = PhysicalStoragePath("user/alice/model", "catalog", "shared-assets") - if public || physical != "user/alice/model" { - t.Fatalf("personal path = %q, %v", physical, public) - } -} - -func TestLogicalStoragePathIsConfigurable(t *testing.T) { - logical, matched := LogicalStoragePath("shared-assets/Models/a/b", "catalog", "shared-assets") - if !matched || logical != "catalog/Models/a/b" { - t.Fatalf("LogicalStoragePath() = %q, %v", logical, matched) - } - logical, matched = LogicalStoragePath("user/alice/model", "catalog", "shared-assets") - if matched || logical != "user/alice/model" { - t.Fatalf("non-public path = %q, %v", logical, matched) - } -} - -func TestReconcilePublicLinksOnePathlessHistoricalDatasetByIdentity(t *testing.T) { - db, err := gorm.Open(sqlite.Open("file:pathless_reconcile?mode=memory&cache=shared"), &gorm.Config{}) - if err != nil { - t.Fatal(err) - } - if err := db.AutoMigrate( - &model.ModelDatasetSource{}, - &model.Dataset{}, - &model.ModelDownload{}, - &model.ModelDatasetDiscovery{}, - ); err != nil { - t.Fatal(err) - } - download := model.ModelDownload{ - Name: "owner/pathless-model", Source: model.ModelSourceHuggingFace, - Category: model.DownloadCategoryModel, Path: "public/Models/owner/pathless-model", - Status: model.ModelDownloadStatusReady, CreatorID: 1, - } - if err := db.Create(&download).Error; err != nil { - t.Fatal(err) - } - dataset := model.Dataset{Name: download.Name, URL: "", Type: model.DataTypeModel} - if err := db.Create(&dataset).Error; err != nil { - t.Fatal(err) - } - - report, err := ReconcilePublic(context.Background(), db, nil, &ReconcileOptions{ - Apply: true, LogicalPublicPrefix: "public", PhysicalPublicPrefix: "shared", - MaxReadmeBytes: 1024, Now: time.Now(), - }) - if err != nil { - t.Fatalf("ReconcilePublic() error = %v", err) - } - if report.DatasetLinks != 1 { - t.Fatalf("report.DatasetLinks = %d", report.DatasetLinks) - } - if report.PathlessDatasets != 1 { - t.Fatalf("report.PathlessDatasets = %d", report.PathlessDatasets) - } - if err := db.First(&dataset, dataset.ID).Error; err != nil { - t.Fatal(err) - } - if dataset.ModelDatasetSourceID == nil { - t.Fatal("pathless historical dataset was not linked to its unique source identity") - } - var discovery model.ModelDatasetDiscovery - if err := db.Where("discovery_key = ?", "dataset:"+fmt.Sprint(dataset.ID)).First(&discovery).Error; err != nil { - t.Fatalf("pathless discovery was not recorded: %v", err) - } - if discovery.Path != "" || discovery.Status != model.ModelDatasetDiscoveryStatusPathMissing { - t.Fatalf("pathless discovery = %#v", discovery) - } -} - -//nolint:gocyclo // The integration assertion covers inferred source, README, ownership, and discovery linkage. -func TestReconcilePublicCreatesOnlyHighConfidenceSourceAndRecordsUniqueOwner(t *testing.T) { - db, err := gorm.Open(sqlite.Open("file:provenance_reconcile?mode=memory&cache=shared"), &gorm.Config{}) - if err != nil { - t.Fatal(err) - } - if err := db.AutoMigrate( - &model.ModelDatasetSource{}, &model.Dataset{}, &model.ModelDownload{}, - &model.ModelDatasetDiscovery{}, &model.User{}, - ); err != nil { - t.Fatal(err) - } - uid := "10042" - user := model.User{ - Name: "historical-owner", Role: model.RoleUser, Status: model.StatusActive, - Space: "user/historical-owner", Attributes: datatypes.NewJSONType(model.UserAttribute{UID: &uid}), - } - if err := db.Create(&user).Error; err != nil { - t.Fatal(err) - } - directory := t.TempDir() - mustWriteFile(t, filepath.Join(directory, "README.md"), "# local legacy model") - candidate := Candidate{ - Path: "shared/LLM/falcon-40b-instruct", AbsolutePath: directory, - Type: model.DataTypeModel, Name: "falcon-40b-instruct", - Evidence: model.ModelDatasetDiscoveryEvidence{ - HasConfig: true, WeightFiles: 1, FilesystemUID: uid, - Provider: model.ModelDatasetProviderHuggingFace, - RepositoryID: "tiiuae/falcon-40b-instruct", RepositoryURL: "https://huggingface.co/tiiuae/falcon-40b-instruct", - ProvenanceSource: "git_remote", ProvenanceConfidence: provenanceConfidenceHigh, - }, - } - report, err := ReconcilePublic(context.Background(), db, []Candidate{candidate}, &ReconcileOptions{ - Apply: true, LogicalPublicPrefix: "public", PhysicalPublicPrefix: "shared", - MaxReadmeBytes: 1024, Now: time.Now(), - }) - if err != nil { - t.Fatalf("ReconcilePublic() error = %v", err) - } - if report.InferredSources != 1 || report.OwnerMatches != 1 || report.ProvenanceHighConfidence != 1 { - t.Fatalf("report = %#v", report) - } - var source model.ModelDatasetSource - if err := db.Where("repository_id = ?", "tiiuae/falcon-40b-instruct").First(&source).Error; err != nil { - t.Fatalf("inferred source not created: %v", err) - } - if source.Readme != "# local legacy model" { - t.Fatalf("source.Readme = %q", source.Readme) - } - var discovery model.ModelDatasetDiscovery - if err := db.Where("path = ?", candidate.Path).First(&discovery).Error; err != nil { - t.Fatal(err) - } - evidence := discovery.Evidence.Data() - if discovery.SourceID == nil || *discovery.SourceID != source.ID || - evidence.OwnerUserID == nil || *evidence.OwnerUserID != user.ID || - evidence.OwnerUsername != user.Name { - t.Fatalf("discovery = %#v, evidence = %#v", discovery, evidence) - } -} - -func TestEnrichCandidateOwnersLeavesDuplicateUIDUnassigned(t *testing.T) { - db, err := gorm.Open(sqlite.Open("file:duplicate_uid?mode=memory&cache=shared"), &gorm.Config{}) - if err != nil { - t.Fatal(err) - } - if err := db.AutoMigrate(&model.User{}); err != nil { - t.Fatal(err) - } - uid := "10042" - for index, name := range []string{"first-owner", "second-owner"} { - user := model.User{ - Name: name, Role: model.RoleUser, Status: model.StatusActive, - Space: fmt.Sprintf("user/%d", index), - Attributes: datatypes.NewJSONType(model.UserAttribute{UID: &uid}), - } - if err := db.Create(&user).Error; err != nil { - t.Fatal(err) - } - } - candidates := []Candidate{{Evidence: model.ModelDatasetDiscoveryEvidence{FilesystemUID: uid}}} - matches, err := enrichCandidateOwners(context.Background(), db, candidates) - if err != nil { - t.Fatal(err) - } - if matches != 0 || candidates[0].Evidence.OwnerUserID != nil || candidates[0].Evidence.OwnerUsername != "" { - t.Fatalf("ambiguous UID was assigned: %#v", candidates[0].Evidence) - } -} diff --git a/backend/internal/governance/modeldataset/scanner.go b/backend/internal/governance/modeldataset/scanner.go deleted file mode 100644 index a10203bc4..000000000 --- a/backend/internal/governance/modeldataset/scanner.go +++ /dev/null @@ -1,363 +0,0 @@ -// Copyright 2026 The Crater Project Team, RAIDS-Lab -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package modeldataset - -import ( - "context" - "errors" - "fmt" - "io/fs" - "os" - "path/filepath" - "sort" - "strings" - - "github.com/raids-lab/crater/dao/model" -) - -const maxEvidenceFiles = 16 - -type ScanOptions struct { - StorageRoot string - PublicPrefix string - ModelsSubdirectory string - ModelsSubdirectories []string - DatasetsSubdirectory string - MaxDepth int - ExcludedDirectories []string - WeightPatterns []string - DatasetMarkerPatterns []string -} - -type Candidate struct { - Path string - AbsolutePath string - Type model.DataType - Name string - Evidence model.ModelDatasetDiscoveryEvidence - SizeBytes int64 -} - -func ScanPublic(ctx context.Context, options *ScanOptions) ([]Candidate, error) { - if options.StorageRoot == "" || options.PublicPrefix == "" { - return nil, errors.New("storage root and public prefix are required") - } - if options.MaxDepth < 1 { - return nil, errors.New("max depth must be positive") - } - - publicRoot := filepath.Join(options.StorageRoot, filepath.FromSlash(options.PublicPrefix)) - excluded := makeSet(options.ExcludedDirectories) - candidates := make([]Candidate, 0) - seenPaths := make(map[string]struct{}) - modelRoots, err := modelSubdirectories(options) - if err != nil { - return nil, err - } - for _, subdirectory := range modelRoots { - models, err := scanModels( - ctx, - filepath.Join(publicRoot, filepath.FromSlash(subdirectory)), - options.PublicPrefix, - subdirectory, - options.MaxDepth, - excluded, - options.WeightPatterns, - ) - if err != nil { - return nil, err - } - for index := range models { - candidate := &models[index] - if _, exists := seenPaths[candidate.Path]; exists { - continue - } - seenPaths[candidate.Path] = struct{}{} - candidates = append(candidates, *candidate) - } - } - - if len(options.DatasetMarkerPatterns) > 0 { - datasets, scanErr := scanDatasets( - ctx, - filepath.Join(publicRoot, filepath.FromSlash(options.DatasetsSubdirectory)), - options.PublicPrefix, - options.DatasetsSubdirectory, - options.MaxDepth, - excluded, - options.DatasetMarkerPatterns, - ) - if scanErr != nil { - return nil, scanErr - } - candidates = append(candidates, datasets...) - } - - sort.Slice(candidates, func(i, j int) bool { return candidates[i].Path < candidates[j].Path }) - return candidates, nil -} - -//nolint:gocyclo // Filesystem traversal keeps exclusion, depth, and candidate checks adjacent. -func scanModels( - ctx context.Context, - root, publicPrefix, subdirectory string, - maxDepth int, - excluded map[string]struct{}, - weightPatterns []string, -) ([]Candidate, error) { - if len(weightPatterns) == 0 { - return nil, errors.New("at least one model weight pattern is required") - } - if _, err := os.Stat(root); errors.Is(err, os.ErrNotExist) { - return []Candidate{}, nil - } else if err != nil { - return nil, err - } - - seen := make(map[string]struct{}) - candidates := make([]Candidate, 0) - err := filepath.WalkDir(root, func(path string, entry fs.DirEntry, walkErr error) error { - if walkErr != nil { - return walkErr - } - if err := ctx.Err(); err != nil { - return err - } - depth, err := relativeDepth(root, path) - if err != nil { - return err - } - if entry.IsDir() { - if path != root && isExcluded(entry.Name(), excluded) { - return filepath.SkipDir - } - if depth > maxDepth { - return filepath.SkipDir - } - return nil - } - if depth > maxDepth || entry.Name() != "config.json" { - return nil - } - - directory := filepath.Dir(path) - if _, ok := seen[directory]; ok { - return nil - } - candidate, ok, err := modelCandidate(directory, root, publicPrefix, subdirectory, weightPatterns) - if err != nil { - return err - } - if ok { - seen[directory] = struct{}{} - candidates = append(candidates, candidate) - } - return nil - }) - return candidates, err -} - -func modelCandidate( - directory, root, publicPrefix, subdirectory string, - weightPatterns []string, -) (Candidate, bool, error) { - entries, err := os.ReadDir(directory) - if err != nil { - return Candidate{}, false, err - } - evidence := model.ModelDatasetDiscoveryEvidence{HasConfig: true} - var sizeBytes int64 - for _, entry := range entries { - if entry.IsDir() { - continue - } - name := entry.Name() - if isReadme(name) { - evidence.HasReadme = true - } - if !matchesAny(name, weightPatterns) { - continue - } - evidence.WeightFiles++ - if len(evidence.MatchedFiles) < maxEvidenceFiles { - evidence.MatchedFiles = append(evidence.MatchedFiles, name) - } - info, infoErr := entry.Info() - if infoErr != nil { - return Candidate{}, false, infoErr - } - sizeBytes += info.Size() - } - if evidence.WeightFiles == 0 { - return Candidate{}, false, nil - } - if err := collectFilesystemEvidence(directory, &evidence); err != nil { - return Candidate{}, false, err - } - relative, err := filepath.Rel(root, directory) - if err != nil { - return Candidate{}, false, err - } - return Candidate{ - Path: normalizedPath(publicPrefix, subdirectory, relative), - AbsolutePath: directory, - Type: model.DataTypeModel, - Name: filepath.Base(directory), - Evidence: evidence, - SizeBytes: sizeBytes, - }, true, nil -} - -func modelSubdirectories(options *ScanOptions) ([]string, error) { - configured := options.ModelsSubdirectories - if len(configured) == 0 { - configured = []string{options.ModelsSubdirectory} - } - result := make([]string, 0, len(configured)) - seen := make(map[string]struct{}) - for _, value := range configured { - value = strings.TrimSpace(value) - if filepath.IsAbs(value) || strings.HasPrefix(value, "/") || strings.Contains(value, "\\") { - return nil, fmt.Errorf("model subdirectory must stay below the public prefix: %q", value) - } - value = strings.Trim(filepath.ToSlash(value), "/") - if value == "" || value == "." { - continue - } - if strings.HasPrefix(value, "../") || value == ".." { - return nil, fmt.Errorf("model subdirectory must stay below the public prefix: %q", value) - } - if _, exists := seen[value]; exists { - continue - } - seen[value] = struct{}{} - result = append(result, value) - } - return result, nil -} - -func scanDatasets( - ctx context.Context, - root, publicPrefix, subdirectory string, - maxDepth int, - excluded map[string]struct{}, - markerPatterns []string, -) ([]Candidate, error) { - if _, err := os.Stat(root); errors.Is(err, os.ErrNotExist) { - return []Candidate{}, nil - } else if err != nil { - return nil, err - } - - seen := make(map[string]struct{}) - candidates := make([]Candidate, 0) - err := filepath.WalkDir(root, func(path string, entry fs.DirEntry, walkErr error) error { - if walkErr != nil { - return walkErr - } - if err := ctx.Err(); err != nil { - return err - } - depth, err := relativeDepth(root, path) - if err != nil { - return err - } - if entry.IsDir() { - if path != root && isExcluded(entry.Name(), excluded) { - return filepath.SkipDir - } - if depth > maxDepth { - return filepath.SkipDir - } - return nil - } - if depth > maxDepth || !matchesAny(entry.Name(), markerPatterns) { - return nil - } - directory := filepath.Dir(path) - if _, ok := seen[directory]; ok { - return nil - } - relative, err := filepath.Rel(root, directory) - if err != nil { - return err - } - seen[directory] = struct{}{} - candidates = append(candidates, Candidate{ - Path: normalizedPath(publicPrefix, subdirectory, relative), - AbsolutePath: directory, - Type: model.DataTypeDataset, - Name: filepath.Base(directory), - Evidence: model.ModelDatasetDiscoveryEvidence{ - HasReadme: isReadme(entry.Name()), - MatchedFiles: []string{entry.Name()}, - }, - }) - return nil - }) - return candidates, err -} - -func relativeDepth(root, path string) (int, error) { - relative, err := filepath.Rel(root, path) - if err != nil { - return 0, err - } - if relative == "." { - return 0, nil - } - return len(strings.Split(filepath.ToSlash(relative), "/")), nil -} - -func normalizedPath(parts ...string) string { - cleaned := make([]string, 0, len(parts)) - for _, part := range parts { - if part == "" || part == "." { - continue - } - cleaned = append(cleaned, filepath.ToSlash(part)) - } - return strings.TrimPrefix(filepath.ToSlash(filepath.Clean(filepath.Join(cleaned...))), "/") -} - -func makeSet(values []string) map[string]struct{} { - result := make(map[string]struct{}, len(values)) - for _, value := range values { - value = strings.TrimSpace(value) - if value != "" { - result[value] = struct{}{} - } - } - return result -} - -func isExcluded(name string, excluded map[string]struct{}) bool { - _, ok := excluded[name] - return ok -} - -func matchesAny(name string, patterns []string) bool { - for _, pattern := range patterns { - matched, err := filepath.Match(pattern, name) - if err == nil && matched { - return true - } - } - return false -} - -func isReadme(name string) bool { - return strings.EqualFold(name, "README") || strings.EqualFold(name, "README.md") -} diff --git a/backend/internal/governance/modeldataset/scanner_test.go b/backend/internal/governance/modeldataset/scanner_test.go deleted file mode 100644 index d990a2ed6..000000000 --- a/backend/internal/governance/modeldataset/scanner_test.go +++ /dev/null @@ -1,208 +0,0 @@ -// Copyright 2026 The Crater Project Team, RAIDS-Lab -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package modeldataset - -import ( - "context" - "os" - "path/filepath" - "testing" - - "github.com/raids-lab/crater/dao/model" -) - -func TestScanPublicFindsCompleteModelsAndSkipsExcludedTrees(t *testing.T) { - root := t.TempDir() - modelDir := filepath.Join(root, "shared", "Models", "owner", "model", "huggingface", "default") - mustWriteFile(t, filepath.Join(modelDir, "config.json"), "{}") - mustWriteFile(t, filepath.Join(modelDir, "README.md"), "# model") - mustWriteFile(t, filepath.Join(modelDir, "model-00001-of-00002.safetensors"), "weights") - incompleteDir := filepath.Join(root, "shared", "Models", "owner", "incomplete") - mustWriteFile(t, filepath.Join(incompleteDir, "config.json"), "{}") - excludedDir := filepath.Join(root, "shared", "Models", "tests", "fixture") - mustWriteFile(t, filepath.Join(excludedDir, "config.json"), "{}") - mustWriteFile(t, filepath.Join(excludedDir, "model.safetensors"), "weights") - - candidates, err := ScanPublic(context.Background(), &ScanOptions{ - StorageRoot: root, - PublicPrefix: "shared", - ModelsSubdirectory: "Models", - DatasetsSubdirectory: "Datasets", - MaxDepth: 8, - ExcludedDirectories: []string{"tests"}, - WeightPatterns: []string{"*.safetensors"}, - }) - if err != nil { - t.Fatalf("ScanPublic() error = %v", err) - } - if len(candidates) != 1 { - t.Fatalf("ScanPublic() candidates = %#v, want one", candidates) - } - candidate := candidates[0] - if candidate.Path != "shared/Models/owner/model/huggingface/default" { - t.Fatalf("candidate.Path = %q", candidate.Path) - } - if candidate.Type != model.DataTypeModel || !candidate.Evidence.HasReadme || candidate.Evidence.WeightFiles != 1 { - t.Fatalf("candidate = %#v", candidate) - } -} - -func TestScanPublicDatasetDiscoveryIsOptIn(t *testing.T) { - root := t.TempDir() - datasetDir := filepath.Join(root, "shared", "Datasets", "owner", "dataset") - mustWriteFile(t, filepath.Join(datasetDir, "dataset_info.json"), "{}") - - withoutMarkers, err := ScanPublic(context.Background(), &ScanOptions{ - StorageRoot: root, - PublicPrefix: "shared", - ModelsSubdirectory: "Models", - DatasetsSubdirectory: "Datasets", - MaxDepth: 8, - WeightPatterns: []string{"*.safetensors"}, - }) - if err != nil { - t.Fatalf("ScanPublic() without markers error = %v", err) - } - if len(withoutMarkers) != 0 { - t.Fatalf("dataset discovery must be disabled without explicit markers: %#v", withoutMarkers) - } - - withMarkers, err := ScanPublic(context.Background(), &ScanOptions{ - StorageRoot: root, - PublicPrefix: "shared", - ModelsSubdirectory: "Models", - DatasetsSubdirectory: "Datasets", - MaxDepth: 8, - WeightPatterns: []string{"*.safetensors"}, - DatasetMarkerPatterns: []string{"dataset_info.json"}, - }) - if err != nil { - t.Fatalf("ScanPublic() with markers error = %v", err) - } - if len(withMarkers) != 1 || withMarkers[0].Type != model.DataTypeDataset { - t.Fatalf("dataset candidates = %#v", withMarkers) - } -} - -func TestScanPublicSupportsMultipleModelRootsAndGitProvenance(t *testing.T) { - root := t.TempDir() - legacyDir := filepath.Join(root, "shared", "LLM", "falcon-40b-instruct") - mustWriteFile(t, filepath.Join(legacyDir, "config.json"), `{}`) - mustWriteFile(t, filepath.Join(legacyDir, "model.safetensors"), "weights") - mustWriteFile(t, filepath.Join(legacyDir, ".git", "config"), "url = git@hf.co:tiiuae/falcon-40b-instruct\n") - managedDir := filepath.Join(root, "shared", "Models", "owner", "managed") - mustWriteFile(t, filepath.Join(managedDir, "config.json"), `{}`) - mustWriteFile(t, filepath.Join(managedDir, "model.safetensors"), "weights") - - candidates, err := ScanPublic(context.Background(), &ScanOptions{ - StorageRoot: root, PublicPrefix: "shared", - ModelsSubdirectory: "ignored-when-list-is-set", - ModelsSubdirectories: []string{"Models", "LLM", "LLM"}, - MaxDepth: 8, WeightPatterns: []string{"*.safetensors"}, - }) - if err != nil { - t.Fatalf("ScanPublic() error = %v", err) - } - if len(candidates) != 2 { - t.Fatalf("candidates = %#v", candidates) - } - legacy := candidates[0] - if legacy.Path != "shared/LLM/falcon-40b-instruct" { - t.Fatalf("legacy.Path = %q", legacy.Path) - } - if legacy.Evidence.Provider != model.ModelDatasetProviderHuggingFace || - legacy.Evidence.RepositoryID != "tiiuae/falcon-40b-instruct" || - legacy.Evidence.ProvenanceSource != "git_remote" || - legacy.Evidence.ProvenanceConfidence != provenanceConfidenceHigh { - t.Fatalf("legacy evidence = %#v", legacy.Evidence) - } - if legacy.Evidence.FilesystemUID == "" || legacy.Evidence.FilesystemGID == "" || - legacy.Evidence.ModifiedAt == nil { - t.Fatalf("filesystem evidence = %#v", legacy.Evidence) - } -} - -func TestScanPublicTreatsConfigAndAmbiguousReadmeAsHints(t *testing.T) { - root := t.TempDir() - directory := filepath.Join(root, "shared", "Models", "llama2-7b") - mustWriteFile(t, filepath.Join(directory, "config.json"), `{"_name_or_path":"meta-llama/Llama-2-7b-hf"}`) - mustWriteFile(t, filepath.Join(directory, "model.safetensors"), "weights") - mustWriteFile(t, filepath.Join(directory, "README.md"), - "Derived from https://huggingface.co/other/base and https://modelscope.cn/models/other/tokenizer") - - candidates, err := ScanPublic(context.Background(), &ScanOptions{ - StorageRoot: root, PublicPrefix: "shared", ModelsSubdirectory: "Models", - MaxDepth: 8, WeightPatterns: []string{"*.safetensors"}, - }) - if err != nil { - t.Fatalf("ScanPublic() error = %v", err) - } - evidence := candidates[0].Evidence - if evidence.Provider != "" || evidence.RepositoryID != "meta-llama/Llama-2-7b-hf" || - evidence.ProvenanceSource != "config_name_or_path" || - evidence.ProvenanceConfidence != provenanceConfidenceMedium { - t.Fatalf("evidence = %#v", evidence) - } - if len(evidence.CandidateURLs) != 2 { - t.Fatalf("candidate URLs = %#v", evidence.CandidateURLs) - } -} - -func TestScanPublicRecognizesMatchingModelScopeReadme(t *testing.T) { - root := t.TempDir() - directory := filepath.Join(root, "shared", "Models", "Qwen2.5-7B-Instruct") - mustWriteFile(t, filepath.Join(directory, "config.json"), `{}`) - mustWriteFile(t, filepath.Join(directory, "model.safetensors"), "weights") - mustWriteFile(t, filepath.Join(directory, "README.md"), - "Model card: https://modelscope.cn/models/Qwen/Qwen2.5-7B-Instruct") - - candidates, err := ScanPublic(context.Background(), &ScanOptions{ - StorageRoot: root, PublicPrefix: "shared", ModelsSubdirectory: "Models", - MaxDepth: 8, WeightPatterns: []string{"*.safetensors"}, - }) - if err != nil { - t.Fatalf("ScanPublic() error = %v", err) - } - evidence := candidates[0].Evidence - if evidence.Provider != model.ModelDatasetProviderModelScope || - evidence.RepositoryID != "Qwen/Qwen2.5-7B-Instruct" || - evidence.ProvenanceSource != "readme_url" || - evidence.ProvenanceConfidence != provenanceConfidenceHigh { - t.Fatalf("evidence = %#v", evidence) - } -} - -func TestScanPublicRejectsModelRootsOutsidePublicPrefix(t *testing.T) { - for _, modelRoot := range []string{"../private", "/private", `..\private`} { - _, err := ScanPublic(context.Background(), &ScanOptions{ - StorageRoot: t.TempDir(), PublicPrefix: "shared", - ModelsSubdirectories: []string{modelRoot}, MaxDepth: 8, - WeightPatterns: []string{"*.safetensors"}, - }) - if err == nil { - t.Fatalf("ScanPublic() accepted model root %q outside the public prefix", modelRoot) - } - } -} - -func mustWriteFile(t *testing.T, path, content string) { - t.Helper() - if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { - t.Fatal(err) - } - if err := os.WriteFile(path, []byte(content), 0o600); err != nil { - t.Fatal(err) - } -} diff --git a/backend/internal/governance/modeldataset/source_logo.go b/backend/internal/governance/modeldataset/source_logo.go deleted file mode 100644 index 55546380a..000000000 --- a/backend/internal/governance/modeldataset/source_logo.go +++ /dev/null @@ -1,188 +0,0 @@ -// Copyright 2026 The Crater Project Team, RAIDS-Lab -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package modeldataset - -import ( - "context" - "encoding/json" - "errors" - "fmt" - stdhtml "html" - "io" - "net/http" - "net/url" - "regexp" - "strings" -) - -const ( - maxLogoRedirects = 5 - maxModelScopePageBytes = 2 * 1024 * 1024 - maxAvatarResponseBytes = 1024 * 1024 -) - -var modelScopeAvatarPattern = regexp.MustCompile( - `https://resou(?:r)?ces\.modelscope\.cn/avatar/[A-Za-z0-9._%/-]+`, -) - -// FetchHuggingFaceAvatarURL resolves an organization or user avatar from the -// same official metadata endpoints used by the batch metadata refresher. -func FetchHuggingFaceAvatarURL( - ctx context.Context, client *http.Client, baseEndpoints []string, owner string, -) (string, error) { - escapedOwner := url.PathEscape(owner) - var lastErr error - for _, baseEndpoint := range baseEndpoints { - endpoints := []string{ - strings.TrimRight(baseEndpoint, "/") + "/api/organizations/" + escapedOwner + "/overview", - strings.TrimRight(baseEndpoint, "/") + "/api/users/" + escapedOwner + "/overview", - } - for _, endpoint := range endpoints { - response, err := getSourceResponse(ctx, client, endpoint) - if err != nil { - lastErr = err - continue - } - if response.StatusCode == http.StatusNotFound { - response.Body.Close() - continue - } - if response.StatusCode != http.StatusOK { - lastErr = fmt.Errorf("source returned HTTP %d", response.StatusCode) - response.Body.Close() - continue - } - var payload struct { - AvatarURL string `json:"avatarUrl"` - } - decodeErr := json.NewDecoder(io.LimitReader(response.Body, maxAvatarResponseBytes)).Decode(&payload) - response.Body.Close() - if decodeErr != nil { - lastErr = decodeErr - continue - } - return payload.AvatarURL, nil - } - } - return "", lastErr -} - -// FetchModelScopeAvatarURL extracts the organization avatar rendered in a -// ModelScope repository page because its OpenAPI response currently omits it. -func FetchModelScopeAvatarURL( - ctx context.Context, client *http.Client, repositoryURL string, -) (string, error) { - response, err := getSourceResponse(ctx, client, repositoryURL) - if err != nil { - return "", err - } - defer response.Body.Close() - if response.StatusCode != http.StatusOK { - return "", fmt.Errorf("source returned HTTP %d", response.StatusCode) - } - page, err := io.ReadAll(io.LimitReader(response.Body, maxModelScopePageBytes+1)) - if err != nil { - return "", err - } - if len(page) > maxModelScopePageBytes { - return "", errors.New("ModelScope repository page is too large") - } - normalized := stdhtml.UnescapeString(strings.ReplaceAll(string(page), `\u002F`, "/")) - normalized = strings.ReplaceAll(normalized, `\/`, "/") - avatarURL := modelScopeAvatarPattern.FindString(normalized) - if avatarURL == "" { - return "", nil - } - parsedAvatarURL, err := url.Parse(avatarURL) - if err != nil { - return "", err - } - query := parsedAvatarURL.Query() - query.Set("x-oss-process", "image/resize,m_lfit,w_128,h_128") - parsedAvatarURL.RawQuery = query.Encode() - return parsedAvatarURL.String(), nil -} - -// FetchSourceLogo downloads and validates a source logo. Every redirect is -// checked against the exact host allowlist to prevent metadata-driven SSRF. -func FetchSourceLogo( - ctx context.Context, - client *http.Client, - endpoint string, - allowedHosts []string, - maxBytes int64, -) (data []byte, contentType string, err error) { - if err := ValidateSourceLogoURL(endpoint, allowedHosts); err != nil { - return nil, "", err - } - logoClient := *client - logoClient.CheckRedirect = func(request *http.Request, via []*http.Request) error { - if len(via) >= maxLogoRedirects { - return errors.New("logo source exceeded redirect limit") - } - return ValidateSourceLogoURL(request.URL.String(), allowedHosts) - } - response, err := getSourceResponse(ctx, &logoClient, endpoint) - if err != nil { - return nil, "", err - } - defer response.Body.Close() - if response.StatusCode != http.StatusOK { - return nil, "", fmt.Errorf("logo source returned HTTP %d", response.StatusCode) - } - contentType = strings.TrimSpace(strings.Split(response.Header.Get("Content-Type"), ";")[0]) - data, err = io.ReadAll(io.LimitReader(response.Body, maxBytes+1)) - if err != nil { - return nil, "", err - } - if int64(len(data)) > maxBytes { - return nil, "", fmt.Errorf("logo exceeds %d bytes", maxBytes) - } - if !strings.HasPrefix(contentType, "image/") { - detectedContentType := strings.TrimSpace(strings.Split(http.DetectContentType(data), ";")[0]) - if !strings.HasPrefix(detectedContentType, "image/") { - return nil, "", fmt.Errorf("logo source returned unsupported Content-Type %q", contentType) - } - contentType = detectedContentType - } - return data, contentType, nil -} - -func ValidateSourceLogoURL(endpoint string, allowedHosts []string) error { - parsed, err := url.Parse(endpoint) - if err != nil { - return fmt.Errorf("invalid logo URL: %w", err) - } - if parsed.Scheme != "https" || parsed.Hostname() == "" || parsed.User != nil { - return errors.New("logo URL must be an absolute HTTPS URL without credentials") - } - - host := strings.ToLower(strings.TrimSuffix(parsed.Hostname(), ".")) - for _, allowedHost := range allowedHosts { - allowedHost = strings.ToLower(strings.TrimSuffix(strings.TrimSpace(allowedHost), ".")) - if host == allowedHost { - return nil - } - } - return fmt.Errorf("logo host %q is not allowed", host) -} - -func getSourceResponse(ctx context.Context, client *http.Client, endpoint string) (*http.Response, error) { - request, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint, http.NoBody) - if err != nil { - return nil, err - } - return client.Do(request) -} diff --git a/backend/internal/handler/agent/agent.go b/backend/internal/handler/agent/agent.go new file mode 100644 index 000000000..db18fc6db --- /dev/null +++ b/backend/internal/handler/agent/agent.go @@ -0,0 +1,144 @@ +package agent + +import ( + "net/http" + "time" + + "github.com/gin-gonic/gin" + "k8s.io/client-go/kubernetes" + "sigs.k8s.io/controller-runtime/pkg/client" + + "github.com/raids-lab/crater/internal/handler" + "github.com/raids-lab/crater/internal/service" + "github.com/raids-lab/crater/pkg/crclient" + "github.com/raids-lab/crater/pkg/monitor" +) + +//nolint:gochecknoinits // Handler managers are registered during package initialization. +func init() { + handler.Registers = append(handler.Registers, NewAgentMgr) +} + +const ( + agentDefaultPythonServiceURL = "http://localhost:8000" + agentChatMessageMaxRunes = 4000 + agentLLMPreflightTimeout = 8 * time.Second + + agentToolStatusSuccess = "success" + agentToolStatusError = "error" + agentToolStatusConfirmationRequired = "confirmation_required" + agentToolStatusAwaitConfirm = "await_confirm" + agentToolStatusAwaitingConfirmation = "awaiting_confirmation" + + agentTurnStatusRunning = "running" + agentTurnStatusCompleted = "completed" + agentTurnStatusFailed = "failed" + agentTurnStatusCancelled = "cancelled" //nolint:misspell // API/frontend status currently uses British spelling. + agentToolStatusRejected = "rejected" + + agentRoleSingleAgent = "single_agent" + agentRoleCoordinator = "coordinator" + agentMessageRoleAssistant = "assistant" + agentSessionSourceUser = "user" + agentSessionSourceChat = "chat" + agentSessionSourceAdmin = "admin" + agentSessionSourceOpsAudit = "ops_audit" + agentSessionSourceSystem = "system" + agentToolAuditSourceBackend = "backend" + + // Read-only tool names + agentToolGetJobDetail = "get_job_detail" + agentToolGetJobEvents = "get_job_events" + agentToolGetJobLogs = "get_job_logs" + agentToolDiagnoseJob = "diagnose_job" + agentToolGetDiagnosticCtx = "get_diagnostic_context" + agentToolSearchSimilarFail = "search_similar_failures" + agentToolQueryJobMetrics = "query_job_metrics" + agentToolAnalyzeQueue = "analyze_queue_status" + agentToolRealtimeCapacity = "get_realtime_capacity" + agentToolListImages = "list_available_images" + agentToolListGPUModels = "list_available_gpu_models" + agentToolCheckQuota = "check_quota" + agentToolListUserJobs = "list_user_jobs" + + // New read-only tools + agentToolGetJobTemplates = "get_job_templates" + agentToolResourceRecommend = "get_resource_recommendation" + + // Write tools that require user confirmation before execution + agentToolResubmitJob = "resubmit_job" + agentToolStopJob = "stop_job" + agentToolDeleteJob = "delete_job" + agentToolCreateJupyter = "create_jupyter_job" + agentToolCreateWebIDE = "create_webide_job" + agentToolCreateCustom = "create_custom_job" + agentToolCreatePytorch = "create_pytorch_job" + agentToolCreateTensorflow = "create_tensorflow_job" +) + +type AgentMgr struct { + name string + client client.Client + kubeClient kubernetes.Interface + nodeClient *crclient.NodeClient + promClient monitor.PrometheusInterface + agentService *service.AgentService + configService *service.ConfigService + jobSubmitter handler.JobMutationSubmitter + jobReader handler.JobInsightReader + imageReader handler.ImageInsightReader + httpClient *http.Client +} + +func NewAgentMgr(conf *handler.RegisterConfig) handler.Manager { + return &AgentMgr{ + name: "agent", + client: conf.Client, + kubeClient: conf.KubeClient, + nodeClient: &crclient.NodeClient{Client: conf.Client, KubeClient: conf.KubeClient, PrometheusClient: conf.PrometheusClient}, + promClient: conf.PrometheusClient, + agentService: service.NewAgentService(), + configService: conf.ConfigService, + jobSubmitter: handler.NewJobMutationSubmitter(conf), + jobReader: handler.NewJobInsightReader(conf), + imageReader: handler.NewImageInsightReader(conf), + // Do not use Client.Timeout for agent SSE streaming. The Python side may + // legitimately take minutes for multi-agent runs, and the per-request + // context in python_proxy.go provides the actual upper bound. + httpClient: &http.Client{}, + } +} + +func (mgr *AgentMgr) GetName() string { return mgr.name } + +func (mgr *AgentMgr) RegisterPublic(g *gin.RouterGroup) { + g.POST("/tools/execute", mgr.ExecuteTool) +} + +func (mgr *AgentMgr) RegisterProtected(g *gin.RouterGroup) { + g.POST("/chat", mgr.Chat) + g.POST("/ask/stream", mgr.Ask) + g.POST("/chat/confirm", mgr.ConfirmToolExecution) + g.POST("/chat/resume", mgr.ResumeAfterConfirmation) + g.GET("/config-summary", mgr.GetAgentConfigSummary) + g.GET("/sessions", mgr.ListSessions) + g.PUT("/sessions/:sessionId/pin", mgr.UpdateSessionPin) + g.PUT("/sessions/:sessionId/title", mgr.UpdateSessionTitle) + g.DELETE("/sessions/:sessionId", mgr.DeleteSession) + g.GET("/sessions/:sessionId/messages", mgr.GetSessionMessages) + g.GET("/sessions/:sessionId/tool-calls", mgr.GetSessionToolCalls) + g.GET("/sessions/:sessionId/turns", mgr.GetSessionTurns) + g.GET("/turns/:turnId/events", mgr.GetTurnEvents) + g.POST("/chat/parameter-update", mgr.HandleParameterUpdate) + + // Feedback + g.PUT("/feedbacks", mgr.UpsertFeedback) + g.POST("/feedbacks/submit", mgr.SubmitFeedback) + g.POST("/feedbacks/quick-submit", mgr.QuickSubmitFeedback) + g.PUT("/feedbacks/enrich", mgr.EnrichFeedback) + g.GET("/feedbacks", mgr.ListFeedbacks) +} + +func (mgr *AgentMgr) RegisterInternal(_ *gin.RouterGroup) {} + +func (mgr *AgentMgr) RegisterAdmin(_ *gin.RouterGroup) {} diff --git a/backend/internal/handler/agent/capabilities.go b/backend/internal/handler/agent/capabilities.go new file mode 100644 index 000000000..2ad85fbed --- /dev/null +++ b/backend/internal/handler/agent/capabilities.go @@ -0,0 +1,144 @@ +package agent + +import ( + "sort" + + "github.com/raids-lab/crater/internal/util" +) + +var agentUserTools = []string{ + agentToolAnalyzeQueue, + agentToolCheckQuota, + agentToolCreateCustom, + agentToolCreateJupyter, + agentToolCreatePytorch, + agentToolCreateTensorflow, + agentToolCreateWebIDE, + agentToolDeleteJob, + agentToolDiagnoseJob, + agentToolGetDiagnosticCtx, + agentToolGetJobDetail, + agentToolGetJobEvents, + agentToolGetJobLogs, + agentToolGetJobTemplates, + agentToolListGPUModels, + agentToolListImages, + agentToolListUserJobs, + agentToolQueryJobMetrics, + agentToolRealtimeCapacity, + agentToolResubmitJob, + agentToolResourceRecommend, + agentToolSearchSimilarFail, + agentToolStopJob, +} + +var agentConfirmToolSet = map[string]struct{}{ + agentToolCreateCustom: {}, + agentToolCreateJupyter: {}, + agentToolCreatePytorch: {}, + agentToolCreateTensorflow: {}, + agentToolCreateWebIDE: {}, + agentToolDeleteJob: {}, + agentToolResubmitJob: {}, + agentToolStopJob: {}, +} + +//nolint:gocyclo // Tool descriptions are intentionally centralized for capability export. +func agentToolCompactDescription(toolName string) string { + switch toolName { + case agentToolGetJobDetail: + return "读取作业状态、资源、时间线和终止信息" + case agentToolGetJobEvents: + return "读取作业相关 Kubernetes 事件" + case agentToolGetJobLogs: + return "读取作业日志尾部或按关键词过滤" + case agentToolDiagnoseJob: + return "执行规则诊断并输出故障分类和根因" + case agentToolGetDiagnosticCtx: + return "读取完整诊断上下文" + case agentToolSearchSimilarFail: + return "检索相似历史失败案例" + case agentToolQueryJobMetrics: + return "读取 GPU、CPU、内存等监控指标" + case agentToolAnalyzeQueue: + return "分析 Pending 或排队原因" + case agentToolRealtimeCapacity: + return "读取实时资源容量概览" + case agentToolListImages: + return "列出当前可见镜像" + case agentToolListGPUModels: + return "列出当前可用 GPU 型号和数量" + case agentToolCheckQuota: + return "查看账户配额使用情况" + case agentToolListUserJobs: + return "列出当前用户近期作业" + case agentToolGetJobTemplates: + return "列出平台提供的作业模板" + case agentToolResourceRecommend: + return "根据任务描述推荐 CPU/GPU/内存配置" + case agentToolCreateCustom: + return "创建自定义作业，需要确认" + case agentToolCreateJupyter: + return "创建 Jupyter 作业，需要确认" + case agentToolCreatePytorch: + return "创建 PyTorch 分布式作业，需要确认" + case agentToolCreateTensorflow: + return "创建 TensorFlow 分布式作业，需要确认" + case agentToolCreateWebIDE: + return "创建 WebIDE 作业，需要确认" + case agentToolDeleteJob: + return "删除作业，需要确认" + case agentToolResubmitJob: + return "重新提交已有作业，需要确认" + case agentToolStopJob: + return "停止作业，需要确认" + default: + return "平台工具" + } +} + +func buildAgentToolCatalog(enabledTools []string) []map[string]any { + catalog := make([]map[string]any, 0, len(enabledTools)) + for _, toolName := range enabledTools { + mode := "read_only" + if isAgentConfirmTool(toolName) { + mode = "confirm" + } + catalog = append(catalog, map[string]any{ + "name": toolName, + "mode": mode, + "description": agentToolCompactDescription(toolName), + }) + } + return catalog +} + +func (mgr *AgentMgr) buildAgentCapabilities(token util.JWTMessage, page map[string]any) map[string]any { + return buildAgentCapabilitiesWithCatalog(token, page) +} + +func buildAgentCapabilitiesWithCatalog(token util.JWTMessage, page map[string]any) map[string]any { + pageRoute, _ := page["route"].(string) + pageURL, _ := page["url"].(string) + pageScope := agentPageScopeForToken(token, page) + enabledTools := append([]string(nil), agentUserTools...) + sort.Strings(enabledTools) + + confirmTools := make([]string, 0, len(enabledTools)) + for _, name := range enabledTools { + if _, ok := agentConfirmToolSet[name]; ok { + confirmTools = append(confirmTools, name) + } + } + + return map[string]any{ + "enabled_tools": enabledTools, + "confirm_tools": confirmTools, + "tool_catalog": buildAgentToolCatalog(enabledTools), + "surface": map[string]any{ + "page_scope": pageScope, + "page_route": pageRoute, + "page_url": pageURL, + }, + } +} diff --git a/backend/internal/handler/agent/confirmation.go b/backend/internal/handler/agent/confirmation.go new file mode 100644 index 000000000..056eacff6 --- /dev/null +++ b/backend/internal/handler/agent/confirmation.go @@ -0,0 +1,175 @@ +package agent + +import ( + "encoding/json" + "fmt" + "strings" + + "github.com/raids-lab/crater/internal/util" +) + +func (mgr *AgentMgr) buildToolConfirmation(_ util.JWTMessage, toolName string, rawArgs json.RawMessage) AgentToolConfirmation { + confirmation := AgentToolConfirmation{ + ToolName: toolName, + RiskLevel: "high", + Interaction: "approval", + Description: mgr.buildConfirmationDescription(toolName, rawArgs), + PermissionExplanation: buildToolPermissionExplanation(toolName), + RiskExplanation: buildToolRiskExplanation(toolName), + AffectedResources: inferToolAffectedResources(rawArgs), + } + switch toolName { + case agentToolCreateJupyter, agentToolCreateWebIDE, agentToolCreateCustom, + agentToolCreatePytorch, agentToolCreateTensorflow, agentToolResubmitJob: + confirmation.Interaction = "form" + confirmation.Form = buildJobForm(toolName, rawArgs) + } + return confirmation +} + +func buildToolPermissionExplanation(toolName string) string { + switch toolName { + case agentToolCreateJupyter, agentToolCreateWebIDE, agentToolCreateCustom, + agentToolCreatePytorch, agentToolCreateTensorflow: + return "需要使用你的作业创建权限提交新的工作负载，并申请表单中的资源。" + case agentToolResubmitJob: + return "需要读取原作业配置，并基于确认后的表单重新提交一个新作业。" + case agentToolStopJob, agentToolDeleteJob: + return "需要确认目标作业属于你或当前身份有权管理该作业。" + default: + return "这是一个需要显式确认的写操作；系统会在你确认后才以当前登录身份执行。" + } +} + +func buildToolRiskExplanation(toolName string) string { + switch toolName { + case agentToolDeleteJob: + return "删除作业会停止相关工作负载并清除记录，请确认目标作业无误。" + case agentToolStopJob: + return "停止作业会中断当前运行任务，但不会创建新作业。" + case agentToolResubmitJob: + return "重提会创建新作业并重新申请资源；原作业不会被自动删除。" + case agentToolCreateJupyter, agentToolCreateWebIDE, agentToolCreateCustom, + agentToolCreatePytorch, agentToolCreateTensorflow: + return "创建作业会占用账户配额和集群资源，配置错误可能导致排队或启动失败。" + default: + return "确认后系统会执行该写操作。" + } +} + +func inferToolAffectedResources(rawArgs json.RawMessage) []string { + args := parseToolArgsMap(rawArgs) + resources := make([]string, 0, 2) + for _, key := range []string{"job_name", "jobName", "name"} { + value := strings.TrimSpace(getToolArgString(args, key, "")) + if value != "" { + resources = append(resources, fmt.Sprintf("作业: %s", value)) + break + } + } + return resources +} + +func (mgr *AgentMgr) buildConfirmationDescription(toolName string, rawArgs json.RawMessage) string { + args := parseToolArgsMap(rawArgs) + switch toolName { + case agentToolDeleteJob: + return fmt.Sprintf("删除作业 %s", getToolArgString(args, "job_name", getToolArgString(args, "jobName", ""))) + case agentToolStopJob: + return fmt.Sprintf("停止作业 %s", getToolArgString(args, "job_name", getToolArgString(args, "jobName", ""))) + case agentToolResubmitJob: + return fmt.Sprintf("重新提交作业 %s", getToolArgString(args, "job_name", getToolArgString(args, "jobName", ""))) + case agentToolCreateJupyter: + return "创建 Jupyter 作业" + case agentToolCreateWebIDE: + return "创建 WebIDE 作业" + case agentToolCreateCustom: + return "创建自定义作业" + case agentToolCreatePytorch: + return "创建 PyTorch 作业" + case agentToolCreateTensorflow: + return "创建 TensorFlow 作业" + default: + return toolName + } +} + +func buildJobForm(toolName string, rawArgs json.RawMessage) *AgentToolForm { + args := parseToolArgsMap(rawArgs) + fields := []AgentToolField{ + { + Key: "name", + Label: "作业名称", + Type: "text", + Required: toolName != agentToolResubmitJob, + DefaultValue: firstArg(args, "name", "job_name", "jobName"), + }, + { + Key: "image_link", + Label: "镜像", + Type: "text", + Required: toolName != agentToolResubmitJob, + DefaultValue: firstArg(args, "image_link", "image", "imageLink"), + }, + { + Key: "cpu", + Label: "CPU", + Type: "text", + Required: false, + DefaultValue: firstArg(args, "cpu"), + }, + { + Key: "memory", + Label: "内存", + Type: "text", + Required: false, + DefaultValue: firstArg(args, "memory"), + }, + { + Key: "gpu_count", + Label: "GPU 数量", + Type: "number", + Required: false, + DefaultValue: firstArg(args, "gpu_count", "gpu"), + }, + { + Key: "gpu_model", + Label: "GPU 型号", + Type: "select", + Required: false, + DefaultValue: firstArg(args, "gpu_model", "gpuModel"), + Options: []AgentToolFieldOption{ + {Label: "不指定", Value: ""}, + {Label: "V100", Value: "v100"}, + {Label: "A100", Value: "a100"}, + {Label: "H100", Value: "h100"}, + {Label: "L40S", Value: "l40s"}, + {Label: "RTX4090", Value: "rtx4090"}, + }, + }, + } + if toolName == agentToolCreateCustom || toolName == agentToolCreatePytorch || toolName == agentToolCreateTensorflow { + fields = append(fields, AgentToolField{ + Key: "command", + Label: "启动命令", + Type: "textarea", + Required: toolName == agentToolCreateCustom, + DefaultValue: firstArg(args, "command"), + }) + } + return &AgentToolForm{ + Title: "确认作业配置", + Description: "可在执行前调整关键作业参数。", + Fields: fields, + SubmitLabel: "确认执行", + } +} + +func firstArg(args map[string]any, keys ...string) string { + for _, key := range keys { + if value := strings.TrimSpace(getToolArgString(args, key, "")); value != "" { + return value + } + } + return "" +} diff --git a/backend/internal/handler/agent/continuation.go b/backend/internal/handler/agent/continuation.go new file mode 100644 index 000000000..a8cfc7edb --- /dev/null +++ b/backend/internal/handler/agent/continuation.go @@ -0,0 +1,529 @@ +package agent + +import ( + "context" + "encoding/json" + "fmt" + "sort" + "strings" + + "github.com/raids-lab/crater/dao/model" +) + +const ( + agentMissingToolRank = 1 << 30 + agentUnrankedToolRank = 1 << 20 + agentSourceEventContentLimit = 1200 +) + +func (mgr *AgentMgr) buildAgentContinuation(ctx context.Context, sessionID string) map[string]any { + if mgr == nil || mgr.agentService == nil || sessionID == "" { + return nil + } + + turns, err := mgr.agentService.ListTurns(ctx, sessionID) + if err != nil || len(turns) == 0 { + return nil + } + + latestTurn := turns[0] + continuation := map[string]any{ + "source_turn_id": latestTurn.TurnID, + "source_turn_status": latestTurn.Status, + } + meaningful := false + + if clarification := mgr.loadClarificationContinuation(ctx, latestTurn); len(clarification) > 0 { + continuation["clarification"] = clarification + meaningful = true + } + if pendingConfirmations := mgr.loadPendingConfirmationContinuation(ctx, latestTurn); len(pendingConfirmations) > 0 { + continuation["pending_confirmations"] = pendingConfirmations + meaningful = true + } + + if !meaningful { + return nil + } + return continuation +} + +func (mgr *AgentMgr) buildResumeContinuation( + ctx context.Context, + turn *model.AgentTurn, + toolCall *model.AgentToolCall, +) map[string]any { + if mgr == nil || mgr.agentService == nil || turn == nil || toolCall == nil { + return nil + } + + continuation := map[string]any{ + "source_turn_id": turn.TurnID, + "source_turn_status": turn.Status, + } + + resume := mgr.buildToolCallResumeResult(turn, toolCall) + if confirmationResults := mgr.loadConfirmationResultContinuations(ctx, turn); len(confirmationResults) > 0 { + resume["confirmation_results"] = confirmationResults + } + if workflow := mgr.loadWorkflowState(ctx, turn); len(workflow) > 0 { + continuation["workflow"] = workflow + resume["workflow"] = workflow + } + if sourceTurnContext := mgr.loadSourceTurnContext(ctx, turn); len(sourceTurnContext) > 0 { + continuation["source_turn_context"] = sourceTurnContext + resume["source_turn_context"] = sourceTurnContext + if originalUserMessage := historyStringValue(sourceTurnContext["original_user_message"]); strings.TrimSpace(originalUserMessage) != "" { + continuation["original_user_message"] = originalUserMessage + resume["original_user_message"] = originalUserMessage + } + } + continuation["resume_after_confirmation"] = resume + return continuation +} + +func (mgr *AgentMgr) buildToolCallResumeResult(turn *model.AgentTurn, toolCall *model.AgentToolCall) map[string]any { + if mgr == nil || turn == nil || toolCall == nil { + return nil + } + result := map[string]any{ + "confirm_id": fmt.Sprintf("%d", toolCall.ID), + "tool_name": toolCall.ToolName, + "action_title": mgr.buildConfirmationDescription(toolCall.ToolName, json.RawMessage(toolCall.ToolArgs)), + "action_intent": agentActionIntentFromToolName(toolCall.ToolName), + "tool_args": parseToolArgsMap(json.RawMessage(toolCall.ToolArgs)), + "result_status": toolCall.ResultStatus, + "confirmed": toolCall.UserConfirmed != nil && *toolCall.UserConfirmed, + "source_turn_id": turn.TurnID, + } + if parsedResult, _ := parseToolCallResult(toolCall); parsedResult != nil { + result["result"] = parsedResult + } + return result +} + +func (mgr *AgentMgr) loadConfirmationResultContinuations(ctx context.Context, turn *model.AgentTurn) []map[string]any { + if mgr == nil || mgr.agentService == nil || turn == nil || turn.TurnID == "" { + return nil + } + toolCalls, err := mgr.agentService.ListToolCallsByTurn(ctx, turn.TurnID) + if err != nil { + return nil + } + results := make([]map[string]any, 0) + orderedToolCalls := mgr.orderToolCallsForWorkflow(toolCalls, mgr.loadWorkflowState(ctx, turn)) + for _, toolCall := range orderedToolCalls { + if toolCall == nil || toolCall.ResultStatus == agentToolStatusAwaitConfirm { + continue + } + if toolCall.UserConfirmed == nil { + continue + } + if resumeResult := mgr.buildToolCallResumeResult(turn, toolCall); len(resumeResult) > 0 { + results = append(results, resumeResult) + } + } + return results +} + +func (mgr *AgentMgr) loadClarificationContinuation(ctx context.Context, turn *model.AgentTurn) map[string]any { + if mgr == nil || mgr.agentService == nil || turn == nil || turn.TurnID == "" { + return nil + } + + events, err := mgr.agentService.ListRunEvents(ctx, turn.TurnID) + if err != nil { + return nil + } + for i := len(events) - 1; i >= 0; i-- { + event := events[i] + if event == nil || event.EventType != "final_answer" || len(event.Metadata) == 0 { + continue + } + metadata := parseAgentJSONMap(json.RawMessage(event.Metadata)) + if continuation, ok := metadata["continuation"].(map[string]any); ok && len(continuation) > 0 { + return continuation + } + } + return nil +} + +func (mgr *AgentMgr) loadPendingConfirmationContinuation(ctx context.Context, turn *model.AgentTurn) []map[string]any { + if mgr == nil || mgr.agentService == nil || turn == nil || + turn.Status != agentToolStatusAwaitingConfirmation || turn.TurnID == "" { + return nil + } + + toolCalls, err := mgr.agentService.ListToolCallsByTurn(ctx, turn.TurnID) + if err != nil { + return nil + } + results := make([]map[string]any, 0) + workflow := mgr.loadWorkflowState(ctx, turn) + for _, toolCall := range mgr.orderToolCallsForWorkflow(toolCalls, workflow) { + if toolCall == nil || toolCall.ResultStatus != agentToolStatusAwaitConfirm { + continue + } + result := map[string]any{ + "confirm_id": fmt.Sprintf("%d", toolCall.ID), + "tool_name": toolCall.ToolName, + "action_intent": agentActionIntentFromToolName(toolCall.ToolName), + "tool_args": parseToolArgsMap(json.RawMessage(toolCall.ToolArgs)), + "result_status": toolCall.ResultStatus, + "agent_role": toolCall.AgentRole, + "source_turn_id": turn.TurnID, + } + if len(workflow) > 0 { + result["workflow"] = workflow + } + results = append(results, result) + } + return results +} + +func (mgr *AgentMgr) hasOtherPendingConfirmations( + ctx context.Context, + turnID string, + excludeToolCallID uint, +) bool { + if mgr == nil || mgr.agentService == nil || turnID == "" { + return false + } + toolCalls, err := mgr.agentService.ListToolCallsByTurn(ctx, turnID) + if err != nil { + return false + } + for _, toolCall := range toolCalls { + if toolCall == nil || toolCall.ID == excludeToolCallID { + continue + } + if toolCall.ResultStatus == agentToolStatusAwaitConfirm { + return true + } + } + return false +} + +//nolint:gocyclo // Ordering matches pending confirmations against several IDs/signatures from Python workflow state. +func (mgr *AgentMgr) orderToolCallsForWorkflow( + toolCalls []*model.AgentToolCall, + workflow map[string]any, +) []*model.AgentToolCall { + if len(toolCalls) <= 1 { + return toolCalls + } + ordered := append([]*model.AgentToolCall(nil), toolCalls...) + rankByConfirmID := map[string]int{} + rankByToolCallID := map[string]int{} + rankBySignature := map[string]int{} + + addRank := func(index int, item any) { + switch typed := item.(type) { + case string: + if typed != "" { + rankByConfirmID[typed] = index + } + case map[string]any: + confirmID := historyStringValue(typed["confirm_id"]) + if confirmID != "" { + rankByConfirmID[confirmID] = index + } + toolCallID := historyStringValue(typed["tool_call_id"]) + if toolCallID == "" { + toolCallID = historyStringValue(typed["toolCallId"]) + } + if toolCallID != "" { + rankByToolCallID[toolCallID] = index + } + toolName := historyStringValue(typed["tool_name"]) + if toolName == "" { + toolName = historyStringValue(typed["toolName"]) + } + args := typed["tool_args"] + if args == nil { + args = typed["toolArgs"] + } + if sig := toolCallSignatureFromValues(toolName, args); sig != "" { + rankBySignature[sig] = index + } + } + } + + if values, ok := workflow["pending_confirmations"].([]any); ok { + for index, item := range values { + if itemMap, ok := item.(map[string]any); ok { + if confirmation, ok := itemMap["confirmation"].(map[string]any); ok { + addRank(index, confirmation) + toolName := historyStringValue(confirmation["tool_name"]) + if toolName == "" { + toolName = historyStringValue(confirmation["toolName"]) + } + args := itemMap["tool_args"] + if args == nil { + args = itemMap["toolArgs"] + } + if sig := toolCallSignatureFromValues(toolName, args); sig != "" { + rankBySignature[sig] = index + } + continue + } + } + addRank(index, item) + } + } + if values, ok := workflow["pending_confirmation_ids"].([]any); ok { + for index, item := range values { + addRank(index, item) + } + } + if values, ok := workflow["actions"].([]any); ok { + for index, item := range values { + addRank(index, item) + } + } + + rankFor := func(toolCall *model.AgentToolCall) int { + if toolCall == nil { + return agentMissingToolRank + } + if rank, ok := rankByConfirmID[fmt.Sprintf("%d", toolCall.ID)]; ok { + return rank + } + if rank, ok := rankByToolCallID[toolCall.ToolCallID]; ok { + return rank + } + if sig := toolCallSignatureFromValues(toolCall.ToolName, parseToolArgsMap(json.RawMessage(toolCall.ToolArgs))); sig != "" { + if rank, ok := rankBySignature[sig]; ok { + return rank + } + } + return agentUnrankedToolRank + } + + sort.SliceStable(ordered, func(i, j int) bool { + leftRank := rankFor(ordered[i]) + rightRank := rankFor(ordered[j]) + if leftRank != rightRank { + return leftRank < rightRank + } + if !ordered[i].CreatedAt.Equal(ordered[j].CreatedAt) { + return ordered[i].CreatedAt.Before(ordered[j].CreatedAt) + } + return ordered[i].ID < ordered[j].ID + }) + return ordered +} + +func toolCallSignatureFromValues(toolName string, args any) string { + toolName = historyStringValue(toolName) + if toolName == "" { + return "" + } + normalizedArgs := "{}" + if args != nil { + if marshaled, err := json.Marshal(args); err == nil { + normalizedArgs = string(marshaled) + } + } + return toolName + ":" + normalizedArgs +} + +//nolint:gocyclo // Context reconstruction merges messages, tool calls and recent run events for continuation prompts. +func (mgr *AgentMgr) loadSourceTurnContext(ctx context.Context, turn *model.AgentTurn) map[string]any { + if mgr == nil || mgr.agentService == nil || turn == nil || turn.TurnID == "" { + return nil + } + originalUserMessage := mgr.loadTurnOriginalUserMessage(ctx, turn) + toolCallSnapshots := mgr.loadSourceTurnToolCalls(ctx, turn) + events, err := mgr.agentService.ListRunEvents(ctx, turn.TurnID) + if err != nil || len(events) == 0 { + if originalUserMessage == "" && len(toolCallSnapshots) == 0 { + return nil + } + result := map[string]any{ + "turn_id": turn.TurnID, + "source_turn_status": turn.Status, + "original_user_message": originalUserMessage, + } + if len(toolCallSnapshots) > 0 { + result["tool_calls"] = toolCallSnapshots + } + return result + } + const maxEvents = 80 + start := 0 + if len(events) > maxEvents { + start = len(events) - maxEvents + } + snapshots := make([]map[string]any, 0, len(events)-start) + for _, event := range events[start:] { + if event == nil { + continue + } + item := map[string]any{ + "sequence": event.Sequence, + "event_type": event.EventType, + "status": event.EventStatus, + "agent_id": event.AgentID, + "agent_role": event.AgentRole, + "title": event.Title, + "content": truncateAgentHistoryContent(event.Content, agentSourceEventContentLimit), + "created_at": event.CreatedAt, + } + metadata := parseAgentJSONMap(json.RawMessage(event.Metadata)) + if len(metadata) > 0 { + item["metadata"] = metadata + } + snapshots = append(snapshots, item) + } + if len(snapshots) == 0 { + if originalUserMessage == "" && len(toolCallSnapshots) == 0 { + return nil + } + result := map[string]any{ + "turn_id": turn.TurnID, + "source_turn_status": turn.Status, + "original_user_message": originalUserMessage, + } + if len(toolCallSnapshots) > 0 { + result["tool_calls"] = toolCallSnapshots + } + return result + } + result := map[string]any{ + "turn_id": turn.TurnID, + "source_turn_status": turn.Status, + "original_user_message": originalUserMessage, + "events": snapshots, + } + if len(toolCallSnapshots) > 0 { + result["tool_calls"] = toolCallSnapshots + } + return result +} + +func (mgr *AgentMgr) loadSourceTurnToolCalls(ctx context.Context, turn *model.AgentTurn) []map[string]any { + if mgr == nil || mgr.agentService == nil || turn == nil || turn.TurnID == "" { + return nil + } + toolCalls, err := mgr.agentService.ListToolCallsByTurn(ctx, turn.TurnID) + if err != nil || len(toolCalls) == 0 { + return nil + } + workflow := mgr.loadWorkflowState(ctx, turn) + orderedToolCalls := mgr.orderToolCallsForWorkflow(toolCalls, workflow) + results := make([]map[string]any, 0, len(orderedToolCalls)) + for _, toolCall := range orderedToolCalls { + if toolCall == nil { + continue + } + item := map[string]any{ + "id": fmt.Sprintf("%d", toolCall.ID), + "tool_call_id": toolCall.ToolCallID, + "agent_id": toolCall.AgentID, + "agent_role": toolCall.AgentRole, + "tool_name": toolCall.ToolName, + "tool_args": parseToolArgsMap(json.RawMessage(toolCall.ToolArgs)), + "result_status": toolCall.ResultStatus, + "confirmed": toolCall.UserConfirmed != nil && *toolCall.UserConfirmed, + "created_at": toolCall.CreatedAt, + } + if parsedResult, _ := parseToolCallResult(toolCall); parsedResult != nil { + item["result"] = parsedResult + } + results = append(results, item) + } + return results +} + +//nolint:gocyclo // Request matching needs several fallbacks for older session records. +func (mgr *AgentMgr) loadTurnOriginalUserMessage(ctx context.Context, turn *model.AgentTurn) string { + if mgr == nil || mgr.agentService == nil || turn == nil || turn.SessionID == "" { + return "" + } + messages, err := mgr.agentService.ListMessages(ctx, turn.SessionID) + if err != nil { + return "" + } + if turn.RequestID != "" { + for _, msg := range messages { + if msg == nil || msg.Role != "user" { + continue + } + if agentMessageRequestID(msg) == turn.RequestID { + return strings.TrimSpace(msg.Content) + } + } + } + var latest string + for _, msg := range messages { + if msg == nil || msg.Role != "user" { + continue + } + if !msg.CreatedAt.After(turn.StartedAt) || msg.CreatedAt.Equal(turn.StartedAt) { + latest = strings.TrimSpace(msg.Content) + } + } + return latest +} + +func (mgr *AgentMgr) loadWorkflowState(ctx context.Context, turn *model.AgentTurn) map[string]any { + if mgr == nil || mgr.agentService == nil || turn == nil || turn.TurnID == "" { + return nil + } + + events, err := mgr.agentService.ListRunEvents(ctx, turn.TurnID) + if err != nil { + return nil + } + for i := len(events) - 1; i >= 0; i-- { + event := events[i] + if event == nil || len(event.Metadata) == 0 { + continue + } + metadata := parseAgentJSONMap(json.RawMessage(event.Metadata)) + if workflow, ok := metadata["workflow"].(map[string]any); ok && len(workflow) > 0 { + return workflow + } + if continuation, ok := metadata["continuation"].(map[string]any); ok { + if workflow, ok := continuation["workflow"].(map[string]any); ok && len(workflow) > 0 { + return workflow + } + } + } + return nil +} + +func parseAgentJSONMap(raw json.RawMessage) map[string]any { + if len(raw) == 0 { + return nil + } + result := map[string]any{} + if err := json.Unmarshal(raw, &result); err != nil { + return nil + } + return result +} + +func agentActionIntentFromToolName(toolName string) string { + switch toolName { + case agentToolResubmitJob: + return "resubmit" + case agentToolStopJob: + return "stop" + case agentToolDeleteJob: + return "delete" + case agentToolCreateJupyter: + return "create_jupyter_job" + case agentToolCreateWebIDE: + return "create_webide_job" + case agentToolCreateCustom: + return "create_custom_job" + case agentToolCreatePytorch: + return "create_pytorch_job" + case agentToolCreateTensorflow: + return "create_tensorflow_job" + default: + return "" + } +} diff --git a/backend/internal/handler/agent/feedback.go b/backend/internal/handler/agent/feedback.go new file mode 100644 index 000000000..bc83cc98e --- /dev/null +++ b/backend/internal/handler/agent/feedback.go @@ -0,0 +1,235 @@ +package agent + +import ( + "encoding/json" + "errors" + "net/http" + + "github.com/gin-gonic/gin" + "gorm.io/datatypes" + "gorm.io/gorm" + + "github.com/raids-lab/crater/internal/bizerr" + "github.com/raids-lab/crater/internal/resputil" + "github.com/raids-lab/crater/internal/service" + "github.com/raids-lab/crater/internal/util" + + "github.com/raids-lab/crater/dao/model" +) + +// ── Request types ─────────────────────────────────────────────────────────── + +type FeedbackUpsertRequest struct { + SessionID string `json:"sessionId" binding:"required"` + TargetType string `json:"targetType" binding:"required,oneof=message turn"` + TargetID string `json:"targetId" binding:"required"` + Rating int16 `json:"rating" binding:"required,oneof=1 -1"` + Tags json.RawMessage `json:"tags,omitempty"` + Dimensions json.RawMessage `json:"dimensions,omitempty"` + Comment string `json:"comment,omitempty"` +} + +type FeedbackSubmitRequest struct { + SessionID string `json:"sessionId" binding:"required"` + TargetType string `json:"targetType" binding:"required,oneof=message turn"` + TargetID string `json:"targetId" binding:"required"` +} + +// ── Handlers ──────────────────────────────────────────────────────────────── + +// UpsertFeedback godoc +// @Summary Create or update a feedback (draft) +// @Tags agent +// @Accept json +// @Produce json +// @Param request body FeedbackUpsertRequest true "Feedback upsert request" +// @Router /api/v1/agent/feedbacks [put] +func (mgr *AgentMgr) UpsertFeedback(c *gin.Context) { + var req FeedbackUpsertRequest + if err := c.ShouldBindJSON(&req); err != nil { + resputil.HandleError(c, bizerr.BadRequest.ParameterError.Wrap(err, "invalid feedback upsert request")) + return + } + + token := util.GetToken(c) + + fb := &model.AgentFeedback{ + SessionID: req.SessionID, + UserID: token.UserID, + AccountID: token.AccountID, + TargetType: req.TargetType, + TargetID: req.TargetID, + Rating: req.Rating, + Comment: req.Comment, + } + if len(req.Tags) > 0 { + fb.Tags = datatypes.JSON(req.Tags) + } + if len(req.Dimensions) > 0 { + fb.Dimensions = datatypes.JSON(req.Dimensions) + } + + result, created, err := mgr.agentService.UpsertFeedback(c.Request.Context(), fb) + if err != nil { + if errors.Is(err, service.ErrFeedbackAlreadySubmitted) { + resputil.HandleError(c, bizerr.Conflict.ResourceStatusError.New("feedback already submitted")) + return + } + resputil.HandleError(c, bizerr.Internal.DatabaseError.Wrap(err, "failed to upsert feedback")) + return + } + + if created { + c.JSON(http.StatusCreated, gin.H{"code": 201, "data": result}) + } else { + resputil.Success(c, result) + } +} + +// SubmitFeedback godoc +// @Summary Submit a draft feedback (makes it immutable) +// @Tags agent +// @Accept json +// @Produce json +// @Param request body FeedbackSubmitRequest true "Feedback submit request" +// @Router /api/v1/agent/feedbacks/submit [post] +func (mgr *AgentMgr) SubmitFeedback(c *gin.Context) { + var req FeedbackSubmitRequest + if err := c.ShouldBindJSON(&req); err != nil { + resputil.HandleError(c, bizerr.BadRequest.ParameterError.Wrap(err, "invalid feedback submit request")) + return + } + + token := util.GetToken(c) + + result, err := mgr.agentService.SubmitFeedback(c.Request.Context(), token.UserID, req.TargetType, req.TargetID) + if err != nil { + if errors.Is(err, service.ErrFeedbackAlreadySubmitted) { + resputil.HandleError(c, bizerr.Conflict.ResourceStatusError.New("feedback already submitted")) + return + } + resputil.HandleError(c, bizerr.Internal.DatabaseError.Wrap(err, "failed to submit feedback")) + return + } + + resputil.Success(c, result) +} + +// ListFeedbacks godoc +// @Summary List feedbacks for a session +// @Tags agent +// @Produce json +// @Param sessionId query string true "Session ID" +// @Router /api/v1/agent/feedbacks [get] +func (mgr *AgentMgr) ListFeedbacks(c *gin.Context) { + sessionID := c.Query("sessionId") + if sessionID == "" { + resputil.HandleError(c, bizerr.BadRequest.MissingParameter.New("sessionId is required")) + return + } + + token := util.GetToken(c) + + feedbacks, err := mgr.agentService.ListFeedbacks(c.Request.Context(), sessionID, token.UserID) + if err != nil { + resputil.HandleError(c, bizerr.Internal.DatabaseError.Wrap(err, "failed to list feedbacks")) + return + } + + resputil.Success(c, feedbacks) +} + +type FeedbackQuickSubmitRequest struct { + SessionID string `json:"sessionId" binding:"required"` + TargetType string `json:"targetType" binding:"required,oneof=message turn"` + TargetID string `json:"targetId" binding:"required"` + Rating int16 `json:"rating" binding:"required,oneof=1 -1"` + Tags json.RawMessage `json:"tags,omitempty"` + Dimensions json.RawMessage `json:"dimensions,omitempty"` + Comment string `json:"comment,omitempty"` +} + +// QuickSubmitFeedback godoc +// @Summary Create and immediately submit feedback (single operation, no draft step) +// @Tags agent +// @Accept json +// @Produce json +// @Router /api/v1/agent/feedbacks/quick-submit [post] +func (mgr *AgentMgr) QuickSubmitFeedback(c *gin.Context) { + var req FeedbackQuickSubmitRequest + if err := c.ShouldBindJSON(&req); err != nil { + resputil.HandleError(c, bizerr.BadRequest.ParameterError.Wrap(err, "invalid quick feedback submit request")) + return + } + + token := util.GetToken(c) + + fb := &model.AgentFeedback{ + SessionID: req.SessionID, + UserID: token.UserID, + AccountID: token.AccountID, + TargetType: req.TargetType, + TargetID: req.TargetID, + Rating: req.Rating, + Comment: req.Comment, + } + if len(req.Tags) > 0 { + fb.Tags = datatypes.JSON(req.Tags) + } + if len(req.Dimensions) > 0 { + fb.Dimensions = datatypes.JSON(req.Dimensions) + } + + result, err := mgr.agentService.QuickSubmitFeedback(c.Request.Context(), fb) + if err != nil { + resputil.HandleError(c, bizerr.Internal.DatabaseError.Wrap(err, "failed to quick submit feedback")) + return + } + + resputil.Success(c, result) +} + +type FeedbackEnrichRequest struct { + SessionID string `json:"sessionId" binding:"required"` + TargetType string `json:"targetType" binding:"required,oneof=message turn"` + TargetID string `json:"targetId" binding:"required"` + Tags json.RawMessage `json:"tags,omitempty"` + Dimensions json.RawMessage `json:"dimensions,omitempty"` + Comment string `json:"comment,omitempty"` +} + +// EnrichFeedback godoc +// @Summary Add/update detail fields on any feedback (even submitted). Never changes rating or status. +// @Tags agent +// @Accept json +// @Produce json +// @Router /api/v1/agent/feedbacks/enrich [put] +func (mgr *AgentMgr) EnrichFeedback(c *gin.Context) { + var req FeedbackEnrichRequest + if err := c.ShouldBindJSON(&req); err != nil { + resputil.HandleError(c, bizerr.BadRequest.ParameterError.Wrap(err, "invalid feedback enrich request")) + return + } + + token := util.GetToken(c) + + result, err := mgr.agentService.EnrichFeedback( + c.Request.Context(), + token.UserID, + req.TargetType, + req.TargetID, + datatypes.JSON(req.Tags), + datatypes.JSON(req.Dimensions), + req.Comment, + ) + if err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + resputil.HandleError(c, bizerr.NotFound.DataBaseNotFound.New("feedback not found")) + return + } + resputil.HandleError(c, bizerr.Internal.DatabaseError.Wrap(err, "failed to enrich feedback")) + return + } + + resputil.Success(c, result) +} diff --git a/backend/internal/handler/agent/handlers.go b/backend/internal/handler/agent/handlers.go new file mode 100644 index 000000000..1760212bf --- /dev/null +++ b/backend/internal/handler/agent/handlers.go @@ -0,0 +1,1116 @@ +package agent + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "strconv" + "strings" + "time" + "unicode/utf8" + + "github.com/gin-gonic/gin" + "github.com/google/uuid" + "gorm.io/datatypes" + + "github.com/raids-lab/crater/dao/model" + "github.com/raids-lab/crater/internal/bizerr" + "github.com/raids-lab/crater/internal/resputil" + "github.com/raids-lab/crater/internal/service" + "github.com/raids-lab/crater/internal/util" + "github.com/raids-lab/crater/pkg/prompts" +) + +const ( + agentAskTimeout = 120 * time.Second + agentRecentHistoryLimit = 8 + agentSessionTitleMaxRunes = 32 +) + +func agentBadRequest(c *gin.Context, msg string) { + resputil.HandleError(c, bizerr.BadRequest.ParameterError.New(msg)) +} + +func agentInternalError(c *gin.Context, msg string) { + resputil.HandleError(c, bizerr.Internal.ServiceError.New(msg)) +} + +func agentForbidden(c *gin.Context, msg string) { + resputil.HandleError(c, bizerr.Forbidden.PermissionDenied.New(msg)) +} + +func agentNotFound(c *gin.Context, msg string) { + resputil.HandleError(c, bizerr.NotFound.DataBaseNotFound.New(msg)) +} + +func (mgr *AgentMgr) requireConfiguredLLMAvailable(ctx context.Context, userID uint) (*service.LLMConfig, error) { + if mgr.configService == nil { + return nil, bizerr.Internal.ServiceError.New("LLM 配置服务未初始化") + } + + cfg, err := mgr.configService.GetEffectiveLLMConfig(ctx, userID) + if err != nil { + return nil, bizerr.Internal.ServiceError.Wrap(err, "读取 LLM 配置失败") + } + if cfg == nil || + strings.TrimSpace(cfg.BaseURL) == "" || + strings.TrimSpace(cfg.APIKey) == "" || + strings.TrimSpace(cfg.ModelName) == "" { + return nil, bizerr.BadRequest.ParameterError.New( + "LLM 配置不完整，请先在系统配置中设置 BaseURL、API Key 和 Model", + ) + } + + checkCtx, cancel := context.WithTimeout(ctx, agentLLMPreflightTimeout) + defer cancel() + if err := prompts.CheckLLMAvailable( + mgr.httpClient, + checkCtx, + cfg.GetChatCompletionURL(), + cfg.APIKey, + cfg.ModelName, + ); err != nil { + return nil, bizerr.Internal.ServiceError.Wrap( + err, + "LLM 服务不可用，请检查平台 LLM 配置、API Key、模型名或网络连通性", + ) + } + return cfg, nil +} + +func (mgr *AgentMgr) failAgentTurn(ctx context.Context, turnID, message string) { + errorMetadata, _ := json.Marshal(map[string]any{"errorMessage": message}) + _ = mgr.agentService.UpdateTurnStatus(ctx, turnID, agentTurnStatusFailed, nil, errorMetadata) +} + +// Chat godoc +// @Summary Agent chat (SSE) +// @Description Create or continue an agent chat session; streams SSE events from the Python Agent service. +// @Tags agent +// @Accept json +// @Produce text/event-stream +// @Param request body AgentChatRequest true "Chat request" +// @Router /api/v1/agent/chat [post] +// +//nolint:gocyclo // Chat wires validation, session setup, persistence and SSE streaming in one endpoint. +func (mgr *AgentMgr) Chat(c *gin.Context) { + var req AgentChatRequest + if err := c.ShouldBindJSON(&req); err != nil { + agentBadRequest(c, err.Error()) + return + } + if utf8.RuneCountInString(strings.TrimSpace(req.Message)) > agentChatMessageMaxRunes { + agentBadRequest(c, fmt.Sprintf("message exceeds %d characters", agentChatMessageMaxRunes)) + return + } + + token := util.GetToken(c) + orchestrationMode := normalizeOrchestrationMode(req.OrchestrationMode) + requestPageContext := normalizePageContext(req.PageContext) + requestSurface := agentPageScopeForToken(token, requestPageContext) + + sessionID := req.SessionID + if sessionID == "" { + sessionID = uuid.New().String() + } + + session, created, err := mgr.agentService.GetOrCreateSession( + c.Request.Context(), + sessionID, + token.UserID, + token.AccountID, + buildAgentSessionTitle(req.Message, req.PageContext), + req.PageContext, + ) + if errors.Is(err, service.ErrAgentSessionDeleted) { + sessionID = uuid.New().String() + session, created, err = mgr.agentService.GetOrCreateSession( + c.Request.Context(), + sessionID, + token.UserID, + token.AccountID, + buildAgentSessionTitle(req.Message, req.PageContext), + req.PageContext, + ) + } + if err != nil { + agentInternalError(c, fmt.Sprintf("failed to create/load session: %v", err)) + return + } + if session.UserID != token.UserID || session.AccountID != token.AccountID { + agentForbidden(c, "session not found") + return + } + if !isChatSessionSource(session.Source) { + agentForbidden(c, "session not found") + return + } + if !created && agentSessionSurface(session) != requestSurface { + sessionID = uuid.New().String() + _, _, err = mgr.agentService.GetOrCreateSession( + c.Request.Context(), + sessionID, + token.UserID, + token.AccountID, + buildAgentSessionTitle(req.Message, req.PageContext), + req.PageContext, + ) + if err != nil { + agentInternalError(c, fmt.Sprintf("failed to create scoped session: %v", err)) + return + } + } + _ = mgr.agentService.UpdateSessionOrchestrationMode(c.Request.Context(), sessionID, orchestrationMode) + + historyMessages, historyErr := mgr.agentService.ListMessages(c.Request.Context(), sessionID) + if historyErr != nil { + agentInternalError(c, fmt.Sprintf("failed to load session history: %v", historyErr)) + return + } + historyToolCalls, toolCallErr := mgr.agentService.ListToolCalls(c.Request.Context(), sessionID) + if toolCallErr != nil { + agentInternalError(c, fmt.Sprintf("failed to load session tool calls: %v", toolCallErr)) + return + } + + historyForPrompt := filterHistoryMessagesForRequest(historyMessages, req.RequestID) + + if !historyContainsRequestID(historyMessages, req.RequestID) { + userMsg := &model.AgentMessage{ + SessionID: sessionID, + Role: "user", + Content: req.Message, + CreatedAt: time.Now(), + } + if req.RequestID != "" { + metadata, _ := json.Marshal(map[string]any{ + "requestId": req.RequestID, + }) + userMsg.Metadata = metadata + } + if saveErr := mgr.agentService.SaveMessage(c.Request.Context(), userMsg); saveErr != nil { + _ = saveErr + } + } + + effectiveToken := effectiveAgentSessionToken(session, token) + continuationContext := mgr.buildAgentContinuation(c.Request.Context(), sessionID) + turnID := uuid.New().String() + _, err = mgr.agentService.CreateTurn(c.Request.Context(), &model.AgentTurn{ + TurnID: turnID, + SessionID: sessionID, + RequestID: req.RequestID, + OrchestrationMode: orchestrationMode, + Status: agentTurnStatusRunning, + StartedAt: time.Now(), + Metadata: datatypes.JSON(req.ClientContext), + }) + if err != nil { + agentInternalError(c, fmt.Sprintf("failed to create agent turn: %v", err)) + return + } + if _, err = mgr.requireConfiguredLLMAvailable(c.Request.Context(), token.UserID); err != nil { + mgr.failAgentTurn(context.Background(), turnID, err.Error()) + resputil.HandleError(c, err) + return + } + agentPayload := mgr.buildPythonAgentPayload( + c.Request.Context(), + sessionID, + turnID, + req.Message, + effectiveToken, + requestPageContext, + normalizeClientContext(req.ClientContext), + orchestrationMode, + historyForPrompt, + historyToolCalls, + continuationContext, + token.UserID, + ) + mgr.streamPythonAgentResponse(c, sessionID, turnID, orchestrationMode, agentPayload, true) +} + +// Ask godoc +// @Summary Ask-only chat (SSE) +// @Description Create or continue an agent chat session and answer with the configured LLM without tool execution. +// @Tags agent +// @Accept json +// @Produce text/event-stream +// @Param request body AgentAskRequest true "Ask request" +// @Router /api/v1/agent/ask/stream [post] +// +//nolint:gocyclo // Ask endpoint keeps streaming, persistence and error handling together. +func (mgr *AgentMgr) Ask(c *gin.Context) { + var req AgentAskRequest + if err := c.ShouldBindJSON(&req); err != nil { + agentBadRequest(c, err.Error()) + return + } + if utf8.RuneCountInString(strings.TrimSpace(req.Message)) > agentChatMessageMaxRunes { + agentBadRequest(c, fmt.Sprintf("message exceeds %d characters", agentChatMessageMaxRunes)) + return + } + if mgr.configService == nil { + agentInternalError(c, "LLM 配置服务未初始化") + return + } + + token := util.GetToken(c) + pageContext := normalizePageContext(req.PageContext) + requestSurface := agentPageScopeForToken(token, pageContext) + sessionID := req.SessionID + if sessionID == "" { + sessionID = uuid.New().String() + } + + session, created, err := mgr.agentService.GetOrCreateSession( + c.Request.Context(), + sessionID, + token.UserID, + token.AccountID, + buildAgentSessionTitle(req.Message, req.PageContext), + req.PageContext, + ) + if errors.Is(err, service.ErrAgentSessionDeleted) { + sessionID = uuid.New().String() + session, created, err = mgr.agentService.GetOrCreateSession( + c.Request.Context(), + sessionID, + token.UserID, + token.AccountID, + buildAgentSessionTitle(req.Message, req.PageContext), + req.PageContext, + ) + } + if err != nil { + agentInternalError(c, fmt.Sprintf("failed to create/load session: %v", err)) + return + } + if session.UserID != token.UserID || session.AccountID != token.AccountID || !isChatSessionSource(session.Source) { + agentForbidden(c, "session not found") + return + } + if !created && agentSessionSurface(session) != requestSurface { + sessionID = uuid.New().String() + _, _, err = mgr.agentService.GetOrCreateSession( + c.Request.Context(), + sessionID, + token.UserID, + token.AccountID, + buildAgentSessionTitle(req.Message, req.PageContext), + req.PageContext, + ) + if err != nil { + agentInternalError(c, fmt.Sprintf("failed to create scoped ask session: %v", err)) + return + } + } + + historyMessages, historyErr := mgr.agentService.ListMessages(c.Request.Context(), sessionID) + if historyErr != nil { + agentInternalError(c, fmt.Sprintf("failed to load session history: %v", historyErr)) + return + } + if !historyContainsRequestID(historyMessages, req.RequestID) { + userMsg := &model.AgentMessage{ + SessionID: sessionID, + Role: "user", + Content: req.Message, + CreatedAt: time.Now(), + } + if req.RequestID != "" { + metadata, _ := json.Marshal(map[string]any{ + "requestId": req.RequestID, + "mode": "ask", + }) + userMsg.Metadata = metadata + } + _ = mgr.agentService.SaveMessage(c.Request.Context(), userMsg) + } + + turnID := uuid.New().String() + _, err = mgr.agentService.CreateTurn(c.Request.Context(), &model.AgentTurn{ + TurnID: turnID, + SessionID: sessionID, + RequestID: req.RequestID, + OrchestrationMode: "ask", + Status: agentTurnStatusRunning, + StartedAt: time.Now(), + Metadata: datatypes.JSON(req.ClientContext), + }) + if err != nil { + agentInternalError(c, fmt.Sprintf("failed to create ask turn: %v", err)) + return + } + + mgr.streamAskResponse(c, sessionID, turnID, req.Message, req.RequestID, historyMessages, token.UserID) +} + +func (mgr *AgentMgr) streamAskResponse( + c *gin.Context, + sessionID string, + turnID string, + message string, + requestID string, + historyMessages []*model.AgentMessage, + userID uint, +) { + cfg, err := mgr.requireConfiguredLLMAvailable(c.Request.Context(), userID) + if err != nil { + mgr.failAgentTurn(context.Background(), turnID, err.Error()) + resputil.HandleError(c, err) + return + } + + c.Header("Content-Type", "text/event-stream") + c.Header("Cache-Control", "no-cache") + c.Header("Connection", "keep-alive") + c.Header("X-Accel-Buffering", "no") + c.Header("X-Agent-Session-ID", sessionID) + c.Header("X-Agent-Turn-ID", turnID) + c.Header("X-Agent-Orchestration-Mode", "ask") + c.Status(http.StatusOK) + + agentID := "ask-1" + agentRole := "ask" + startPayload := map[string]any{ + "turnId": turnID, + "sessionId": sessionID, + "agentId": agentID, + "agentRole": agentRole, + "status": agentTurnStatusRunning, + "summary": "ask 正在回答", + } + _ = writeSyntheticSSEEvent(c, "agent_run_started", startPayload) + + systemPrompt := strings.Join([]string{ + "你是 Crater 平台的 ask 助手。", + "你负责回答用户关于平台、作业、排障和使用方式的问题，但不执行写操作、不创建资源、不调用工具。", + "如果用户请求删除、停止、创建、排空节点等操作，请说明 ask 只能解释和建议；需要执行请切换到 agent。", + "回答要简洁、准确；不知道时说明缺少哪些信息，不要编造。", + }, "\n") + userPrompt := buildAskUserPrompt(message, historyMessages) + ctx, cancel := context.WithTimeout(c.Request.Context(), agentAskTimeout) + defer cancel() + + var full strings.Builder + reply, err := prompts.CallLLMTextStream( + mgr.httpClient, + ctx, + cfg.GetChatCompletionURL(), + cfg.APIKey, + cfg.ModelName, + systemPrompt, + userPrompt, + func(delta string) error { + full.WriteString(delta) + return writeSyntheticSSEEvent(c, "message", map[string]any{ + "turnId": turnID, + "sessionId": sessionID, + "agentId": agentID, + "agentRole": agentRole, + "content": delta, + "partial": true, + }) + }, + ) + if err != nil { + if ctx.Err() != nil || c.Request.Context().Err() != nil { + _ = mgr.agentService.UpdateTurnStatus(context.Background(), turnID, agentTurnStatusCancelled, nil, nil) + return + } + _ = mgr.agentService.UpdateTurnStatus(context.Background(), turnID, agentTurnStatusFailed, nil, nil) + _ = writeSyntheticSSEEvent(c, "error", map[string]any{"message": fmt.Sprintf("ask 调用失败：%v", err)}) + _ = writeSyntheticSSEEvent(c, "done", map[string]any{}) + return + } + if strings.TrimSpace(reply) == "" { + reply = full.String() + } + + assistantMsg := &model.AgentMessage{ + SessionID: sessionID, + Role: agentMessageRoleAssistant, + Content: reply, + CreatedAt: time.Now(), + } + if requestID != "" { + metadata, _ := json.Marshal(map[string]any{ + "requestId": requestID, + "mode": "ask", + }) + assistantMsg.Metadata = metadata + } + var finalMessageID *uint + if saveErr := mgr.agentService.SaveMessage(context.Background(), assistantMsg); saveErr == nil { + finalMessageID = &assistantMsg.ID + } + _ = mgr.agentService.UpdateTurnStatus(context.Background(), turnID, agentTurnStatusCompleted, finalMessageID, nil) + _ = writeSyntheticSSEEvent(c, "final_answer", map[string]any{ + "turnId": turnID, + "sessionId": sessionID, + "agentId": agentID, + "agentRole": agentRole, + "content": reply, + "feedbackTargetId": fmt.Sprintf("%d", assistantMsg.ID), + }) + _ = writeSyntheticSSEEvent(c, "done", map[string]any{}) +} + +func buildAskUserPrompt(message string, historyMessages []*model.AgentMessage) string { + recent := make([]string, 0, agentRecentHistoryLimit) + start := len(historyMessages) - agentRecentHistoryLimit + if start < 0 { + start = 0 + } + for _, msg := range historyMessages[start:] { + role := strings.TrimSpace(msg.Role) + if role == "" || strings.TrimSpace(msg.Content) == "" { + continue + } + recent = append(recent, fmt.Sprintf("%s: %s", role, strings.TrimSpace(msg.Content))) + } + if len(recent) == 0 { + return fmt.Sprintf("用户问题：%s", message) + } + return fmt.Sprintf("最近对话：\n%s\n\n用户当前问题：%s", strings.Join(recent, "\n"), message) +} + +func isChatSessionSource(source string) bool { + normalized := strings.TrimSpace(strings.ToLower(source)) + return normalized == "" || normalized == agentSessionSourceChat +} + +func buildAgentSessionTitle(message string, pageContext json.RawMessage) string { + title := strings.TrimSpace(message) + title = strings.Join(strings.Fields(title), " ") + if title == "" { + title = "新的 Agent 会话" + } + + var ctx map[string]any + if len(pageContext) > 0 && json.Unmarshal(pageContext, &ctx) == nil { + if jobName, ok := ctx["jobName"].(string); ok && strings.TrimSpace(jobName) != "" && !strings.Contains(title, jobName) { + title = fmt.Sprintf("%s · %s", strings.TrimSpace(jobName), title) + } else if nodeName, ok := ctx["nodeName"].(string); ok && strings.TrimSpace(nodeName) != "" && !strings.Contains(title, nodeName) { + title = fmt.Sprintf("%s · %s", strings.TrimSpace(nodeName), title) + } + } + + runes := []rune(title) + if len(runes) > agentSessionTitleMaxRunes { + return string(runes[:agentSessionTitleMaxRunes]) + "…" + } + return title +} + +func (mgr *AgentMgr) respondOwnedSessionData( + c *gin.Context, + errorPrefix string, + load func(context.Context, string) (any, error), +) { + sessionID := c.Param("sessionId") + if sessionID == "" { + agentBadRequest(c, "sessionId is required") + return + } + + token := util.GetToken(c) + if _, err := mgr.agentService.GetOwnedSession(c.Request.Context(), sessionID, token.UserID); err != nil { + agentForbidden(c, "session not found") + return + } + + data, err := load(c.Request.Context(), sessionID) + if err != nil { + agentInternalError(c, fmt.Sprintf("%s: %v", errorPrefix, err)) + return + } + resputil.Success(c, data) +} + +// ResumeAfterConfirmation godoc +// @Summary Resume an agent turn after a confirmation result +// @Description Resumes a paused agent turn after the confirmation result has been recorded. +// @Tags agent +// @Accept json +// @Produce text/event-stream +// @Param request body AgentResumeRequest true "Resume request" +// @Router /api/v1/agent/chat/resume [post] +// +//nolint:gocyclo // Resume handles both streamed continuation and already-settled confirmation outcomes. +func (mgr *AgentMgr) ResumeAfterConfirmation(c *gin.Context) { + var req AgentResumeRequest + if err := c.ShouldBindJSON(&req); err != nil { + agentBadRequest(c, err.Error()) + return + } + + token := util.GetToken(c) + confirmID, err := strconv.ParseUint(req.ConfirmID, 10, 64) + if err != nil { + agentBadRequest(c, "invalid confirmId") + return + } + toolCall, err := mgr.agentService.GetToolCallByID(c.Request.Context(), uint(confirmID)) + if err != nil { + agentNotFound(c, "confirmation result not found") + return + } + if toolCall.ResultStatus == agentToolStatusAwaitConfirm { + agentBadRequest(c, "confirmation has not completed yet") + return + } + if mgr.hasOtherPendingConfirmations(c.Request.Context(), toolCall.TurnID, toolCall.ID) { + resputil.Success(c, AgentToolResponse{ + Status: agentToolStatusAwaitingConfirmation, + Message: "仍有其他待确认操作，请处理完本轮所有确认卡后再继续。", + }) + return + } + if toolCall.ResultStatus == agentToolStatusRejected { + turnID := toolCall.TurnID + if turnID == "" { + turnID = uuid.New().String() + } + mgr.streamConfirmationOutcome(c, toolCall.SessionID, turnID, toolCall) + return + } + + session, err := mgr.agentService.GetOwnedSession(c.Request.Context(), toolCall.SessionID, token.UserID) + if err != nil { + agentForbidden(c, "confirmation result not found") + return + } + + sourceTurn, sourceTurnErr := mgr.agentService.GetTurn(c.Request.Context(), toolCall.TurnID) + if sourceTurnErr == nil { + orchestrationMode := normalizeOrchestrationMode(sourceTurn.OrchestrationMode) + sessionToken, tokenErr := mgr.getSessionToken(c.Request.Context(), session) + if tokenErr != nil { + agentInternalError(c, fmt.Sprintf("failed to resolve session actor: %v", tokenErr)) + return + } + historyMessages, historyErr := mgr.agentService.ListMessages(c.Request.Context(), toolCall.SessionID) + if historyErr != nil { + agentInternalError(c, fmt.Sprintf("failed to load session history: %v", historyErr)) + return + } + historyToolCalls, toolCallErr := mgr.agentService.ListToolCalls(c.Request.Context(), toolCall.SessionID) + if toolCallErr != nil { + agentInternalError(c, fmt.Sprintf("failed to load session tool calls: %v", toolCallErr)) + return + } + + if err := mgr.agentService.UpdateTurnStatus(c.Request.Context(), sourceTurn.TurnID, agentTurnStatusRunning, nil, nil); err != nil { + agentInternalError(c, fmt.Sprintf("failed to resume source turn: %v", err)) + return + } + if _, err = mgr.requireConfiguredLLMAvailable(c.Request.Context(), sessionToken.UserID); err != nil { + mgr.failAgentTurn(context.Background(), sourceTurn.TurnID, err.Error()) + resputil.HandleError(c, err) + return + } + + agentPayload := mgr.buildPythonAgentPayload( + c.Request.Context(), + toolCall.SessionID, + sourceTurn.TurnID, + "继续完成上一轮计划", + sessionToken, + normalizePageContext(json.RawMessage(session.PageContext)), + normalizeClientContext(json.RawMessage(sourceTurn.Metadata)), + orchestrationMode, + historyMessages, + historyToolCalls, + mgr.buildResumeContinuation(c.Request.Context(), sourceTurn, toolCall), + sessionToken.UserID, + ) + mgr.streamPythonAgentResponse(c, toolCall.SessionID, sourceTurn.TurnID, orchestrationMode, agentPayload, true) + return + } + + turnID := toolCall.TurnID + if turnID == "" { + turnID = uuid.New().String() + } + mgr.streamConfirmationOutcome(c, toolCall.SessionID, turnID, toolCall) +} + +// ConfirmToolExecution godoc +// @Summary Confirm or reject a write tool operation +// @Description Called by the frontend after the user confirms or rejects a write operation (e.g., stop_job). +// @Tags agent +// @Accept json +// @Produce json +// @Param request body ConfirmToolRequest true "Confirmation request" +// @Success 200 {object} resputil.Response[AgentToolResponse] +// @Router /api/v1/agent/chat/confirm [post] +// +//nolint:gocyclo // Confirmation handling validates ownership, executes/rejects and records the result atomically. +func (mgr *AgentMgr) ConfirmToolExecution(c *gin.Context) { + var req ConfirmToolRequest + if err := c.ShouldBindJSON(&req); err != nil { + agentBadRequest(c, err.Error()) + return + } + + token := util.GetToken(c) + confirmID, err := strconv.ParseUint(req.ConfirmID, 10, 64) + if err != nil { + agentBadRequest(c, "invalid confirmId") + return + } + toolCall, err := mgr.agentService.GetToolCallByID(c.Request.Context(), uint(confirmID)) + if err != nil { + agentNotFound(c, "pending action not found") + return + } + session, err := mgr.agentService.GetOwnedSession(c.Request.Context(), toolCall.SessionID, token.UserID) + if err != nil { + agentForbidden(c, "pending action not found") + return + } + if toolCall.ResultStatus != agentToolStatusAwaitConfirm { + agentBadRequest(c, "pending action is no longer awaiting confirmation") + return + } + + if !req.Confirmed { + summary := mgr.buildToolOutcomeMessage(toolCall.ToolName, agentToolStatusRejected, nil, "Operation rejected by user.") + confirmed := false + if updateErr := mgr.agentService.UpdateToolCallOutcome( + c.Request.Context(), + toolCall.ID, + agentToolStatusRejected, + json.RawMessage(toolCall.ToolResult), + &confirmed, + ); updateErr != nil { + agentInternalError(c, fmt.Sprintf("failed to update pending action: %v", updateErr)) + return + } + resultBytes, _ := json.Marshal(map[string]any{ + "confirmId": req.ConfirmID, + "confirmed": false, + }) + resputil.Success(c, AgentToolResponse{ + Status: agentToolStatusRejected, + Result: resultBytes, + Message: summary, + }) + return + } + + sessionToken, tokenErr := mgr.getSessionToken(c.Request.Context(), session) + if tokenErr != nil { + agentInternalError(c, fmt.Sprintf("failed to resolve session actor: %v", tokenErr)) + return + } + if authErr := authorizeAgentToolForSession(session, sessionToken, toolCall.ToolName); authErr != nil { + agentForbidden(c, authErr.Error()) + return + } + + mergedArgs, mergeErr := mergeToolArgsWithPayload(json.RawMessage(toolCall.ToolArgs), req.Payload) + if mergeErr != nil { + agentBadRequest(c, mergeErr.Error()) + return + } + if len(req.Payload) > 0 && string(req.Payload) != "null" { + if updateErr := mgr.agentService.UpdateToolCallArgs(c.Request.Context(), toolCall.ID, mergedArgs); updateErr != nil { + agentInternalError(c, fmt.Sprintf("failed to persist confirmation payload: %v", updateErr)) + return + } + } + + start := time.Now() + executionBackend := strings.TrimSpace(toolCall.ExecutionBackend) + if executionBackend == "" { + executionBackend = normalizeExecutionBackend(toolCall.ToolName, "") + } + + result, execErr := mgr.executeWriteTool(c, sessionToken, toolCall.ToolName, mergedArgs) + latencyMs := int(time.Since(start).Milliseconds()) + + status := agentToolStatusSuccess + var resultBytes json.RawMessage + var responseMsg string + + if execErr != nil { + status = agentToolStatusError + responseMsg = mgr.buildToolOutcomeMessage(toolCall.ToolName, status, nil, execErr.Error()) + if result != nil { + resultBytes, _ = json.Marshal(result) + } + if len(resultBytes) == 0 { + errJSON, _ := json.Marshal(map[string]string{"error": execErr.Error()}) + resultBytes = errJSON + } + } else { + resultBytes, _ = json.Marshal(result) + responseMsg = mgr.buildToolOutcomeMessage(toolCall.ToolName, status, result, "") + } + + confirmed := true + if updateErr := mgr.agentService.UpdateToolCallOutcome( + c.Request.Context(), + toolCall.ID, + status, + resultBytes, + &confirmed, + ); updateErr != nil { + agentInternalError(c, fmt.Sprintf("failed to update pending action: %v", updateErr)) + return + } + recordAgentMutationOperationLog( + c, + toolCall.ToolName, + mergedArgs, + result, + execErr, + executionBackend, + req.ConfirmID, + ) + + resputil.Success(c, AgentToolResponse{ + Status: status, + Result: resultBytes, + Message: responseMsg, + LatencyMs: latencyMs, + }) +} + +// ListSessions godoc +// @Summary List agent chat sessions for the current user +// @Tags agent +// @Produce json +// @Success 200 {object} resputil.Response[any] +// @Router /api/v1/agent/sessions [get] +func (mgr *AgentMgr) ListSessions(c *gin.Context) { + token := util.GetToken(c) + sessions, err := mgr.agentService.ListSessions(c.Request.Context(), token.UserID) + if err != nil { + agentInternalError(c, fmt.Sprintf("failed to list sessions: %v", err)) + return + } + if rawSurface := strings.TrimSpace(c.Query("surface")); rawSurface != "" { + surface := normalizeAgentSurface(rawSurface) + filtered := make([]*model.AgentSession, 0, len(sessions)) + for _, session := range sessions { + if agentSessionMatchesSurface(session, surface) { + filtered = append(filtered, session) + } + } + sessions = filtered + } + resputil.Success(c, sessions) +} + +// UpdateSessionPin godoc +// @Summary Pin or unpin an agent session +// @Tags agent +// @Accept json +// @Produce json +// @Param sessionId path string true "Session ID (UUID)" +// @Param request body AgentSessionPinRequest true "Pin request" +// @Success 200 {object} resputil.Response[any] +// @Router /api/v1/agent/sessions/{sessionId}/pin [put] +func (mgr *AgentMgr) UpdateSessionPin(c *gin.Context) { + sessionID := c.Param("sessionId") + if sessionID == "" { + agentBadRequest(c, "sessionId is required") + return + } + + var req AgentSessionPinRequest + if err := c.ShouldBindJSON(&req); err != nil { + agentBadRequest(c, err.Error()) + return + } + + token := util.GetToken(c) + if _, err := mgr.agentService.GetOwnedSession(c.Request.Context(), sessionID, token.UserID); err != nil { + agentForbidden(c, "session not found") + return + } + if err := mgr.agentService.UpdateSessionPinned(c.Request.Context(), sessionID, req.Pinned); err != nil { + if errors.Is(err, service.ErrAgentSessionPinningUnavailable) { + agentInternalError(c, "session pinning requires a completed database migration") + return + } + agentInternalError(c, fmt.Sprintf("failed to update session pin: %v", err)) + return + } + session, err := mgr.agentService.GetOwnedSession(c.Request.Context(), sessionID, token.UserID) + if err != nil { + agentNotFound(c, "session not found") + return + } + resputil.Success(c, session) +} + +// UpdateSessionTitle godoc +// @Summary Rename an agent session +// @Tags agent +// @Accept json +// @Produce json +// @Param sessionId path string true "Session ID (UUID)" +// @Param request body AgentSessionTitleRequest true "Rename request" +// @Success 200 {object} resputil.Response[any] +// @Router /api/v1/agent/sessions/{sessionId}/title [put] +func (mgr *AgentMgr) UpdateSessionTitle(c *gin.Context) { + sessionID := c.Param("sessionId") + if sessionID == "" { + agentBadRequest(c, "sessionId is required") + return + } + + var req AgentSessionTitleRequest + if err := c.ShouldBindJSON(&req); err != nil { + agentBadRequest(c, err.Error()) + return + } + + token := util.GetToken(c) + if _, err := mgr.agentService.GetOwnedSession(c.Request.Context(), sessionID, token.UserID); err != nil { + agentForbidden(c, "session not found") + return + } + if err := mgr.agentService.UpdateSessionTitle(c.Request.Context(), sessionID, req.Title); err != nil { + agentBadRequest(c, fmt.Sprintf("failed to update session title: %v", err)) + return + } + session, err := mgr.agentService.GetOwnedSession(c.Request.Context(), sessionID, token.UserID) + if err != nil { + agentNotFound(c, "session not found") + return + } + resputil.Success(c, session) +} + +// DeleteSession godoc +// @Summary Soft delete an agent session +// @Tags agent +// @Produce json +// @Param sessionId path string true "Session ID (UUID)" +// @Success 200 {object} resputil.Response[string] +// @Router /api/v1/agent/sessions/{sessionId} [delete] +func (mgr *AgentMgr) DeleteSession(c *gin.Context) { + sessionID := c.Param("sessionId") + if sessionID == "" { + agentBadRequest(c, "sessionId is required") + return + } + + token := util.GetToken(c) + if _, err := mgr.agentService.GetOwnedSession(c.Request.Context(), sessionID, token.UserID); err != nil { + agentForbidden(c, "session not found") + return + } + if err := mgr.agentService.DeleteSession(c.Request.Context(), sessionID); err != nil { + agentInternalError(c, fmt.Sprintf("failed to delete session: %v", err)) + return + } + resputil.Success(c, "ok") +} + +// GetAgentConfigSummary godoc +// @Summary Get agent configuration summary for the current user +// @Tags agent +// @Produce json +// @Success 200 {object} resputil.Response[AgentConfigSummary] +// @Router /api/v1/agent/config-summary [get] +func (mgr *AgentMgr) GetAgentConfigSummary(c *gin.Context) { + summary := AgentConfigSummary{ + DefaultOrchestrationMode: agentRoleSingleAgent, + AvailableModes: []string{agentRoleSingleAgent}, + } + if mgr.configService != nil { + token := util.GetToken(c) + if llmStatus, err := mgr.configService.GetEffectiveLLMConfigStatus(c.Request.Context(), token.UserID); err == nil && llmStatus != nil { + llmSummary := &AgentLLMConfigSummary{ + Source: llmStatus.Source, + UsingConfig: llmStatus.Complete, + UsingPersonal: llmStatus.UsingPersonal, + Complete: llmStatus.Complete, + HasAPIKey: llmStatus.HasAPIKey, + } + if llmStatus.Config != nil { + llmSummary.BaseURL = strings.TrimSpace(llmStatus.Config.BaseURL) + llmSummary.Model = strings.TrimSpace(llmStatus.Config.ModelName) + } + if !llmStatus.Complete { + llmSummary.FallbackNote = "LLM 配置不完整，Agent 不会使用 crater-agent 本地 llm-clients.json 兜底；请在个人设置或平台设置中配置 BaseURL、Model 和 API Key。" + } + summary.LLM = llmSummary + } + } + agentReq, err := http.NewRequestWithContext( + c.Request.Context(), + http.MethodGet, + mgr.getPythonAgentURL()+"/config-summary", + http.NoBody, + ) + if err == nil { + resp, requestErr := mgr.httpClient.Do(agentReq) + if requestErr == nil { + defer resp.Body.Close() + if resp.StatusCode < http.StatusBadRequest { + var agentSummary AgentConfigSummary + if decodeErr := json.NewDecoder(resp.Body).Decode(&agentSummary); decodeErr == nil { + if len(agentSummary.AvailableModes) > 0 { + summary.AvailableModes = agentSummary.AvailableModes + } + } + } + } + } + resputil.Success(c, summary) +} + +// HandleParameterUpdate godoc +// @Summary Forward a parameter update to the Python Agent service +// @Description Allows the frontend to send parameter adjustments (e.g. form field changes) to the agent mid-session. +// @Tags agent +// @Accept json +// @Produce json +// @Param request body map[string]any true "Parameter update payload" +// @Success 200 {object} resputil.Response[any] +// @Router /api/v1/agent/chat/parameter-update [post] +func (mgr *AgentMgr) HandleParameterUpdate(c *gin.Context) { + var payload map[string]any + if err := c.ShouldBindJSON(&payload); err != nil { + agentBadRequest(c, err.Error()) + return + } + + token := util.GetToken(c) + sessionID, _ := payload["sessionId"].(string) + if sessionID == "" { + agentBadRequest(c, "sessionId is required") + return + } + + if _, err := mgr.agentService.GetOwnedSession(c.Request.Context(), sessionID, token.UserID); err != nil { + agentForbidden(c, "session not found") + return + } + + bodyBytes, err := json.Marshal(payload) + if err != nil { + agentInternalError(c, fmt.Sprintf("failed to marshal payload: %v", err)) + return + } + + agentURL := mgr.getPythonAgentURL() + "/chat/parameter-update" + agentReq, err := http.NewRequestWithContext( + c.Request.Context(), + http.MethodPost, + agentURL, + bytes.NewReader(bodyBytes), + ) + if err != nil { + agentInternalError(c, fmt.Sprintf("failed to create agent request: %v", err)) + return + } + agentReq.Header.Set("Content-Type", "application/json") + if internalToken := mgr.getPythonAgentInternalToken(); internalToken != "" { + agentReq.Header.Set("X-Agent-Internal-Token", internalToken) + } + + resp, err := mgr.httpClient.Do(agentReq) + if err != nil { + agentInternalError(c, fmt.Sprintf("agent service unavailable: %v", err)) + return + } + defer resp.Body.Close() + + var result any + if decodeErr := json.NewDecoder(resp.Body).Decode(&result); decodeErr != nil { + agentInternalError(c, fmt.Sprintf("failed to decode agent response: %v", decodeErr)) + return + } + + if resp.StatusCode >= http.StatusBadRequest { + agentInternalError(c, fmt.Sprintf("agent service returned status %d", resp.StatusCode)) + return + } + + resputil.Success(c, result) +} + +// GetSessionMessages godoc +// @Summary Get messages for a specific agent session +// @Tags agent +// @Produce json +// @Param sessionId path string true "Session ID (UUID)" +// @Success 200 {object} resputil.Response[any] +// @Router /api/v1/agent/sessions/{sessionId}/messages [get] +func (mgr *AgentMgr) GetSessionMessages(c *gin.Context) { + mgr.respondOwnedSessionData(c, "failed to list messages", func(ctx context.Context, sessionID string) (any, error) { + return mgr.agentService.ListMessages(ctx, sessionID) + }) +} + +// GetSessionToolCalls godoc +// @Summary Get tool calls for a specific agent session +// @Tags agent +// @Produce json +// @Param sessionId path string true "Session ID (UUID)" +// @Success 200 {object} resputil.Response[any] +// @Router /api/v1/agent/sessions/{sessionId}/tool-calls [get] +func (mgr *AgentMgr) GetSessionToolCalls(c *gin.Context) { + mgr.respondOwnedSessionData(c, "failed to list tool calls", func(ctx context.Context, sessionID string) (any, error) { + return mgr.agentService.ListToolCalls(ctx, sessionID) + }) +} + +// GetSessionTurns godoc +// @Summary Get turns for a specific agent session +// @Tags agent +// @Produce json +// @Param sessionId path string true "Session ID (UUID)" +// @Success 200 {object} resputil.Response[any] +// @Router /api/v1/agent/sessions/{sessionId}/turns [get] +func (mgr *AgentMgr) GetSessionTurns(c *gin.Context) { + mgr.respondOwnedSessionData(c, "failed to list turns", func(ctx context.Context, sessionID string) (any, error) { + return mgr.agentService.ListTurns(ctx, sessionID) + }) +} + +// GetTurnEvents godoc +// @Summary Get run events for a specific agent turn +// @Tags agent +// @Produce json +// @Param turnId path string true "Turn ID (UUID)" +// @Success 200 {object} resputil.Response[any] +// @Router /api/v1/agent/turns/{turnId}/events [get] +func (mgr *AgentMgr) GetTurnEvents(c *gin.Context) { + turnID := c.Param("turnId") + if turnID == "" { + agentBadRequest(c, "turnId is required") + return + } + + token := util.GetToken(c) + turn, err := mgr.agentService.GetTurn(c.Request.Context(), turnID) + if err != nil { + agentNotFound(c, "turn not found") + return + } + _, err = mgr.agentService.GetOwnedSession(c.Request.Context(), turn.SessionID, token.UserID) + if err != nil { + agentForbidden(c, "turn not found") + return + } + events, err := mgr.agentService.ListRunEvents(c.Request.Context(), turnID) + if err != nil { + agentInternalError(c, fmt.Sprintf("failed to list turn events: %v", err)) + return + } + resputil.Success(c, events) +} diff --git a/backend/internal/handler/agent/job_mutation.go b/backend/internal/handler/agent/job_mutation.go new file mode 100644 index 000000000..0f603d7b6 --- /dev/null +++ b/backend/internal/handler/agent/job_mutation.go @@ -0,0 +1,100 @@ +package agent + +import ( + "encoding/json" + "fmt" + "strings" + + "github.com/gin-gonic/gin" + + "github.com/raids-lab/crater/dao/model" + "github.com/raids-lab/crater/dao/query" + "github.com/raids-lab/crater/internal/bizerr" + "github.com/raids-lab/crater/internal/util" +) + +type agentJobNameArgs struct { + JobName string `json:"job_name"` +} + +func isAgentOwnedJobMutationTool(toolName string) bool { + switch toolName { + case agentToolDeleteJob, agentToolStopJob, agentToolResubmitJob: + return true + default: + return false + } +} + +func agentOwnedJobMutationActionName(toolName string) string { + switch toolName { + case agentToolDeleteJob: + return "删除" + case agentToolStopJob: + return "停止" + case agentToolResubmitJob: + return "重提" + default: + return "操作" + } +} + +func (mgr *AgentMgr) validateOwnedJobMutationBeforeConfirmation( + c *gin.Context, + token util.JWTMessage, + toolName string, + rawArgs json.RawMessage, +) error { + if !isAgentOwnedJobMutationTool(toolName) { + return nil + } + + var args agentJobNameArgs + if err := json.Unmarshal(rawArgs, &args); err != nil { + return bizerr.BadRequest.ParameterError.Wrap(err, "invalid args") + } + args.JobName = strings.TrimSpace(args.JobName) + if args.JobName == "" { + return bizerr.BadRequest.MissingParameter.New("job_name is required") + } + + j := query.Job + jobQuery := j.WithContext(c).Where(j.JobName.Eq(args.JobName)) + if token.RolePlatform != model.RoleAdmin { + jobQuery = jobQuery.Where(j.UserID.Eq(token.UserID), j.AccountID.Eq(token.AccountID)) + } + if _, err := jobQuery.First(); err != nil { + return bizerr.Forbidden.PermissionDenied.New( + fmt.Sprintf("该作业不存在或你没有访问权限，不能发起%s确认", agentOwnedJobMutationActionName(toolName)), + ) + } + return nil +} + +func mergeToolArgsWithPayload(baseArgs, payload json.RawMessage) (json.RawMessage, error) { + if len(payload) == 0 || string(payload) == "null" { + return baseArgs, nil + } + + base := make(map[string]any) + if len(baseArgs) > 0 { + if err := json.Unmarshal(baseArgs, &base); err != nil { + return nil, bizerr.BadRequest.ParameterError.Wrap(err, "invalid stored tool args") + } + } + + incoming := make(map[string]any) + if err := json.Unmarshal(payload, &incoming); err != nil { + return nil, bizerr.BadRequest.ParameterError.Wrap(err, "invalid confirmation payload") + } + + for key, value := range incoming { + base[key] = value + } + + merged, err := json.Marshal(base) + if err != nil { + return nil, bizerr.Internal.ServiceError.Wrap(err, "failed to merge confirmation payload") + } + return merged, nil +} diff --git a/backend/internal/handler/agent/message.go b/backend/internal/handler/agent/message.go new file mode 100644 index 000000000..248481189 --- /dev/null +++ b/backend/internal/handler/agent/message.go @@ -0,0 +1,263 @@ +package agent + +import ( + "encoding/json" + "fmt" + "sort" + "strings" + "time" + + "github.com/raids-lab/crater/dao/model" +) + +const ( + agentHistoryMessageLimit = 24 + agentHistoryContentLimit = 1600 + agentHistoryToolContentLimit = 480 + agentHistoryAssistantMaxChars = 1200 + agentHistoryToolArgsLimit = 180 + agentHistoryToolResultLimit = 220 +) + +type agentHistoryEntry struct { + createdAt time.Time + payload map[string]any +} + +func truncateAgentHistoryContent(content string, maxChars int) string { + content = strings.TrimSpace(content) + if content == "" || maxChars <= 0 { + return content + } + runes := []rune(content) + if len(runes) <= maxChars { + return content + } + return string(runes[:maxChars]) + "..." +} + +func agentMessageRequestID(msg *model.AgentMessage) string { + if msg == nil || len(msg.Metadata) == 0 { + return "" + } + var metadata map[string]any + if err := json.Unmarshal(msg.Metadata, &metadata); err != nil { + return "" + } + requestID, _ := metadata["requestId"].(string) + return requestID +} + +func historyContainsRequestID(messages []*model.AgentMessage, requestID string) bool { + if requestID == "" { + return false + } + for _, msg := range messages { + if agentMessageRequestID(msg) == requestID { + return true + } + } + return false +} + +func filterHistoryMessagesForRequest(messages []*model.AgentMessage, requestID string) []*model.AgentMessage { + if requestID == "" { + return messages + } + filtered := make([]*model.AgentMessage, 0, len(messages)) + for _, msg := range messages { + if agentMessageRequestID(msg) == requestID { + continue + } + filtered = append(filtered, msg) + } + return filtered +} + +func historyStringValue(value any) string { + if value == nil { + return "" + } + s := strings.TrimSpace(fmt.Sprintf("%v", value)) + if s == "" { + return "" + } + return s +} + +func compactHistoryJSON(raw json.RawMessage, maxChars int) string { + if len(raw) == 0 { + return "" + } + + var payload any + if err := json.Unmarshal(raw, &payload); err != nil { + return truncateAgentHistoryContent(string(raw), maxChars) + } + + normalized, err := json.Marshal(payload) + if err != nil { + return truncateAgentHistoryContent(string(raw), maxChars) + } + return truncateAgentHistoryContent(string(normalized), maxChars) +} + +func decorateAssistantHistoryContent(msg *model.AgentMessage, content string) string { + if msg == nil || msg.Role != agentMessageRoleAssistant || len(msg.Metadata) == 0 { + return content + } + + metadata := parseAgentJSONMap(json.RawMessage(msg.Metadata)) + if len(metadata) == 0 { + return content + } + + source := historyStringValue(metadata["source"]) + if source != "tool_confirmation" && source != "confirmation_resume" { + return content + } + + toolName := historyStringValue(metadata["toolName"]) + status := historyStringValue(metadata["status"]) + prefixParts := make([]string, 0, 2) + if toolName != "" { + prefixParts = append(prefixParts, "tool="+toolName) + } + if status != "" { + prefixParts = append(prefixParts, "status="+status) + } + if len(prefixParts) == 0 { + return "【上轮工具结果】" + content + } + return "【上轮工具结果 " + strings.Join(prefixParts, " ") + "】" + content +} + +func buildAgentHistoryMessageEntry(msg *model.AgentMessage) *agentHistoryEntry { + if msg == nil || strings.TrimSpace(msg.Content) == "" { + return nil + } + + maxChars := agentHistoryContentLimit + switch msg.Role { + case "tool": + maxChars = agentHistoryToolContentLimit + case agentMessageRoleAssistant: + maxChars = agentHistoryAssistantMaxChars + } + + content := truncateAgentHistoryContent(msg.Content, maxChars) + if msg.Role == agentMessageRoleAssistant { + content = decorateAssistantHistoryContent(msg, content) + content = truncateAgentHistoryContent(content, maxChars) + } + + payload := map[string]any{ + "role": msg.Role, + "content": content, + } + if msg.Role == "tool" && strings.TrimSpace(msg.ToolCallID) != "" { + payload["tool_call_id"] = msg.ToolCallID + } + + return &agentHistoryEntry{ + createdAt: msg.CreatedAt, + payload: payload, + } +} + +func buildAgentToolCallHistoryContent(toolCall *model.AgentToolCall) string { + if toolCall == nil || strings.TrimSpace(toolCall.ToolName) == "" { + return "" + } + + status := strings.TrimSpace(toolCall.ResultStatus) + parts := []string{ + fmt.Sprintf("tool=%s", toolCall.ToolName), + } + if status != "" { + parts = append(parts, "status="+status) + } + if toolCall.UserConfirmed != nil { + parts = append(parts, fmt.Sprintf("user_confirmed=%t", *toolCall.UserConfirmed)) + } + + if args := compactHistoryJSON(json.RawMessage(toolCall.ToolArgs), agentHistoryToolArgsLimit); args != "" { + parts = append(parts, "args="+args) + } + + switch status { + case agentToolStatusRejected: + parts = append(parts, "result=operation rejected by user") + case agentToolStatusAwaitConfirm, agentToolStatusConfirmationRequired: + parts = append(parts, "result=awaiting user confirmation") + default: + if result := compactHistoryJSON(json.RawMessage(toolCall.ToolResult), agentHistoryToolResultLimit); result != "" { + parts = append(parts, "result="+result) + } + } + + return strings.Join(parts, " ; ") +} + +func buildAgentHistoryToolEntry(toolCall *model.AgentToolCall) *agentHistoryEntry { + if toolCall == nil { + return nil + } + + content := buildAgentToolCallHistoryContent(toolCall) + if strings.TrimSpace(content) == "" { + return nil + } + + toolCallID := strings.TrimSpace(toolCall.ToolCallID) + if toolCallID == "" { + toolCallID = fmt.Sprintf("tool-call-%d", toolCall.ID) + } + + return &agentHistoryEntry{ + createdAt: toolCall.CreatedAt, + payload: map[string]any{ + "role": "tool", + "content": truncateAgentHistoryContent(content, agentHistoryToolContentLimit), + "tool_call_id": toolCallID, + }, + } +} + +func buildAgentHistory(messages []*model.AgentMessage, toolCalls []*model.AgentToolCall) []map[string]any { + if len(messages) == 0 && len(toolCalls) == 0 { + return nil + } + + entries := make([]agentHistoryEntry, 0, len(messages)+len(toolCalls)) + for _, msg := range messages { + if entry := buildAgentHistoryMessageEntry(msg); entry != nil { + entries = append(entries, *entry) + } + } + for _, toolCall := range toolCalls { + if entry := buildAgentHistoryToolEntry(toolCall); entry != nil { + entries = append(entries, *entry) + } + } + if len(entries) == 0 { + return nil + } + + sort.SliceStable(entries, func(i, j int) bool { + return entries[i].createdAt.Before(entries[j].createdAt) + }) + + start := len(entries) - agentHistoryMessageLimit + if start < 0 { + start = 0 + } + history := make([]map[string]any, 0, len(entries)-start) + for _, entry := range entries[start:] { + if len(entry.payload) == 0 { + continue + } + history = append(history, entry.payload) + } + return history +} diff --git a/backend/internal/handler/agent/normalize.go b/backend/internal/handler/agent/normalize.go new file mode 100644 index 000000000..c13acf287 --- /dev/null +++ b/backend/internal/handler/agent/normalize.go @@ -0,0 +1,103 @@ +package agent + +import ( + "encoding/json" + "strconv" + "strings" +) + +func normalizePageContext(raw json.RawMessage) map[string]any { + if len(raw) == 0 { + return map[string]any{} + } + var page map[string]any + if err := json.Unmarshal(raw, &page); err != nil { + return map[string]any{} + } + if jobName, ok := page["jobName"]; ok { + page["job_name"] = jobName + } + if jobStatus, ok := page["jobStatus"]; ok { + page["job_status"] = jobStatus + } + if nodeName, ok := page["nodeName"]; ok { + page["node_name"] = nodeName + } + if entryPoint, ok := page["entryPoint"]; ok { + page["entrypoint"] = entryPoint + } + return page +} + +func normalizeClientContext(raw json.RawMessage) map[string]any { + if len(raw) == 0 { + return map[string]any{} + } + var clientContext map[string]any + if err := json.Unmarshal(raw, &clientContext); err != nil { + return map[string]any{} + } + return clientContext +} + +func normalizeOrchestrationMode(mode string) string { + _ = mode + return agentRoleSingleAgent +} + +func parseToolArgsMap(rawArgs json.RawMessage) map[string]any { + args := map[string]any{} + _ = json.Unmarshal(rawArgs, &args) + return args +} + +func getToolArgString(args map[string]any, key, fallback string) string { + value, _ := args[key].(string) + value = strings.TrimSpace(value) + if value == "" { + return fallback + } + return value +} + +func getToolArgInt(args map[string]any, key string, fallback int) int { + value, ok := args[key] + if !ok || value == nil { + return fallback + } + switch typed := value.(type) { + case float64: + return int(typed) + case int: + return typed + case int32: + return int(typed) + case int64: + return int(typed) + case string: + parsed, err := strconv.Atoi(strings.TrimSpace(typed)) + if err == nil { + return parsed + } + } + return fallback +} + +func getToolArgBool(args map[string]any, key string, fallback bool) bool { + value, ok := args[key] + if !ok || value == nil { + return fallback + } + switch typed := value.(type) { + case bool: + return typed + case string: + switch strings.ToLower(strings.TrimSpace(typed)) { + case "true", "1", "yes", "y", "on": + return true + case "false", "0", "no", "n", "off": + return false + } + } + return fallback +} diff --git a/backend/internal/handler/agent/operation_log_bridge.go b/backend/internal/handler/agent/operation_log_bridge.go new file mode 100644 index 000000000..2ffc225e9 --- /dev/null +++ b/backend/internal/handler/agent/operation_log_bridge.go @@ -0,0 +1,18 @@ +package agent + +import ( + "encoding/json" + + "github.com/gin-gonic/gin" +) + +func recordAgentMutationOperationLog( + _ *gin.Context, + _ string, + _ json.RawMessage, + _ any, + _ error, + _ string, + _ string, +) { +} diff --git a/backend/internal/handler/agent/python_proxy.go b/backend/internal/handler/agent/python_proxy.go new file mode 100644 index 000000000..b24d2e7fa --- /dev/null +++ b/backend/internal/handler/agent/python_proxy.go @@ -0,0 +1,573 @@ +package agent + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "os" + "strings" + "time" + + "github.com/gin-gonic/gin" + "gorm.io/datatypes" + + "github.com/raids-lab/crater/dao/model" + "github.com/raids-lab/crater/dao/query" + "github.com/raids-lab/crater/internal/bizerr" + "github.com/raids-lab/crater/internal/util" + pkgconfig "github.com/raids-lab/crater/pkg/config" +) + +const ( + agentPythonStreamTimeout = 15 * time.Minute + agentSSEInitialBufferSize = 64 * 1024 + agentSSEMaxBufferSize = 1024 * 1024 +) + +// getPythonAgentURL returns the configured Python Agent service base URL. +func (mgr *AgentMgr) getPythonAgentURL() string { + cfg := pkgconfig.GetConfig() + if cfg.Agent.ServiceURL != "" { + return cfg.Agent.ServiceURL + } + return agentDefaultPythonServiceURL +} + +func (mgr *AgentMgr) getPythonAgentInternalToken() string { + cfg := pkgconfig.GetConfig() + if cfg.Agent.InternalToken != "" { + return cfg.Agent.InternalToken + } + return os.Getenv("CRATER_AGENT_INTERNAL_TOKEN") +} + +func (mgr *AgentMgr) isInternalToolRequestAuthorized(c *gin.Context) bool { + internalToken := mgr.getPythonAgentInternalToken() + if internalToken == "" { + return false + } + return c.GetHeader("X-Agent-Internal-Token") == internalToken +} + +func (mgr *AgentMgr) getSessionToken(ctx context.Context, session *model.AgentSession) (util.JWTMessage, error) { + if session == nil || session.UserID == 0 || session.AccountID == 0 { + return util.JWTMessage{}, bizerr.BadRequest.ParameterError.New("invalid session actor") + } + u := query.User + a := query.Account + ua := query.UserAccount + + user, err := u.WithContext(ctx).Where(u.ID.Eq(session.UserID)).First() + if err != nil { + return util.JWTMessage{}, bizerr.Internal.DatabaseError.Wrap(err, "failed to load user") + } + account, err := a.WithContext(ctx).Where(a.ID.Eq(session.AccountID)).First() + if err != nil { + return util.JWTMessage{}, bizerr.Internal.DatabaseError.Wrap(err, "failed to load account") + } + userAccount, err := ua.WithContext(ctx). + Where(ua.UserID.Eq(session.UserID), ua.AccountID.Eq(session.AccountID)). + First() + if err != nil { + return util.JWTMessage{}, bizerr.Internal.DatabaseError.Wrap(err, "failed to load account membership") + } + + token := util.JWTMessage{ + UserID: session.UserID, + AccountID: session.AccountID, + Username: user.Name, + AccountName: account.Name, + RoleAccount: userAccount.Role, + RolePlatform: user.Role, + AccountAccessMode: userAccount.AccessMode, + } + return effectiveAgentSessionToken(session, token), nil +} + +func (mgr *AgentMgr) buildPythonAgentPayload( + ctx context.Context, + sessionID string, + turnID string, + message string, + token util.JWTMessage, + pageContext map[string]any, + clientContext map[string]any, + orchestrationMode string, + historyMessages []*model.AgentMessage, + historyToolCalls []*model.AgentToolCall, + continuation map[string]any, + userID uint, +) AgentTurnRequest { + requestContext := map[string]any{ + "actor": map[string]any{ + "user_id": token.UserID, + "account_id": token.AccountID, + "username": token.Username, + "account_name": token.AccountName, + "role": strings.ToLower(token.RolePlatform.String()), + }, + "page": pageContext, + "client": clientContext, + "history": buildAgentHistory(historyMessages, historyToolCalls), + "continuation": continuation, + "capabilities": mgr.buildAgentCapabilities(token, pageContext), + "orchestration": map[string]any{"mode": normalizeOrchestrationMode(orchestrationMode)}, + } + if llmContext := mgr.buildPythonAgentLLMContext(ctx, userID); len(llmContext) > 0 { + requestContext["llm"] = llmContext + } + + return AgentTurnRequest{ + SessionID: sessionID, + TurnID: turnID, + Message: message, + Context: requestContext, + } +} + +func (mgr *AgentMgr) buildPythonAgentLLMContext(ctx context.Context, userID uint) map[string]any { + source, clientConfig := mgr.buildPythonAgentLLMClientConfig(ctx, userID) + if len(clientConfig) == 0 { + return nil + } + return map[string]any{ + "source": source, + "client_config": clientConfig, + } +} + +func (mgr *AgentMgr) buildPythonAgentLLMClientConfig( + ctx context.Context, + userID uint, +) (source string, clientConfig map[string]any) { + if mgr.configService == nil { + return "", nil + } + status, err := mgr.configService.GetEffectiveLLMConfigStatus(ctx, userID) + if err != nil || status == nil || !status.Complete { + return "", nil + } + clientConfig, err = mgr.configService.GetAgentLLMClientConfigForUser(ctx, userID) + if err != nil || len(clientConfig) == 0 { + return "", nil + } + return status.Source, clientConfig +} + +//nolint:gocyclo,funlen // SSE proxy persists streamed agent events while forwarding them to the browser. +func (mgr *AgentMgr) streamPythonAgentResponse( + c *gin.Context, + sessionID string, + turnID string, + orchestrationMode string, + agentPayload AgentTurnRequest, + persistAssistant bool, +) { + payloadBytes, err := json.Marshal(agentPayload) + if err != nil { + agentInternalError(c, fmt.Sprintf("failed to marshal agent payload: %v", err)) + return + } + + agentURL := mgr.getPythonAgentURL() + "/chat" + // Multi-agent diagnostic runs can legitimately span many LLM/tool hops. + // Keep a bounded request context, but avoid the previous short timeout + // that interrupted SSE before final_answer could be emitted. + agentCtx, cancelAgentReq := context.WithTimeout(context.Background(), agentPythonStreamTimeout) + defer cancelAgentReq() + agentReq, err := http.NewRequestWithContext(agentCtx, http.MethodPost, agentURL, bytes.NewReader(payloadBytes)) + if err != nil { + agentInternalError(c, fmt.Sprintf("failed to create agent request: %v", err)) + return + } + agentReq.Header.Set("Content-Type", "application/json") + agentReq.Header.Set("Accept", "text/event-stream") + + resp, err := mgr.httpClient.Do(agentReq) + if err != nil { + errorMetadata, _ := json.Marshal(map[string]any{ + "errorMessage": fmt.Sprintf("agent service unavailable: %v", err), + }) + _ = mgr.agentService.UpdateTurnStatus(context.Background(), turnID, agentTurnStatusFailed, nil, errorMetadata) + agentInternalError(c, fmt.Sprintf("agent service unavailable: %v", err)) + return + } + defer resp.Body.Close() + if resp.StatusCode >= http.StatusBadRequest { + errorMetadata, _ := json.Marshal(map[string]any{ + "errorMessage": fmt.Sprintf("agent service returned status %d", resp.StatusCode), + }) + _ = mgr.agentService.UpdateTurnStatus(context.Background(), turnID, agentTurnStatusFailed, nil, errorMetadata) + agentInternalError(c, fmt.Sprintf("agent service returned status %d", resp.StatusCode)) + return + } + + c.Header("Content-Type", "text/event-stream") + c.Header("Cache-Control", "no-cache") + c.Header("Connection", "keep-alive") + c.Header("X-Agent-Session-ID", sessionID) + c.Header("X-Agent-Turn-ID", turnID) + c.Header("X-Agent-Orchestration-Mode", orchestrationMode) + + var assistantContent string + turnStatus := agentTurnStatusRunning + var turnMetadata json.RawMessage + eventSequence := 0 + now := func() *time.Time { + timestamp := time.Now() + return ×tamp + } + persistRunEvent := func(eventType string, rawData []byte) { + if strings.TrimSpace(eventType) == "" || len(rawData) == 0 { + return + } + + eventSequence++ + var payloadMap map[string]any + rawContent := strings.TrimSpace(string(rawData)) + if err := json.Unmarshal(rawData, &payloadMap); err != nil { + payloadMap = map[string]any{"content": rawContent} + } + agentID, _ := payloadMap["agentId"].(string) + parentAgentID, _ := payloadMap["parentAgentId"].(string) + agentRole, _ := payloadMap["agentRole"].(string) + eventStatus, _ := payloadMap["status"].(string) + title, _ := payloadMap["title"].(string) + content, _ := payloadMap["content"].(string) + if content == "" { + if summary, _ := payloadMap["summary"].(string); summary != "" { + content = summary + } else if resultSummary, _ := payloadMap["resultSummary"].(string); resultSummary != "" { + content = resultSummary + } + } + if title == "" { + if toolName, _ := payloadMap["toolName"].(string); toolName != "" { + title = toolName + } + } + if agentRole == "" { + agentRole = agentRoleSingleAgent + } + if eventStatus == "" { + switch eventType { + case "tool_call_started", "agent_run_started": + eventStatus = "started" + case "tool_call_confirmation_required": + eventStatus = agentToolStatusAwaitingConfirmation + case agentToolStatusError: + eventStatus = agentToolStatusError + default: + eventStatus = agentTurnStatusCompleted + } + } + metadataBytes, _ := json.Marshal(payloadMap) + _, _ = mgr.agentService.CreateRunEvent(context.Background(), &model.AgentRunEvent{ + TurnID: turnID, + SessionID: sessionID, + AgentID: agentID, + ParentAgentID: parentAgentID, + AgentRole: agentRole, + EventType: eventType, + EventStatus: eventStatus, + Title: title, + Content: content, + Metadata: datatypes.JSON(metadataBytes), + Sequence: eventSequence, + StartedAt: now(), + EndedAt: now(), + }) + + switch eventType { + case "tool_call_confirmation_required": + turnStatus = agentToolStatusAwaitingConfirmation + case agentToolStatusError: + turnStatus = agentTurnStatusFailed + errorMessage := "" + if message, _ := payloadMap["message"].(string); strings.TrimSpace(message) != "" { + errorMessage = message + } else if message, _ := payloadMap["msg"].(string); strings.TrimSpace(message) != "" { + errorMessage = message + } else if strings.TrimSpace(content) != "" { + errorMessage = content + } else if title != "" { + errorMessage = title + } + if strings.TrimSpace(errorMessage) == "" { + errorMessage = "agent execution failed" + } + turnMetadata, _ = json.Marshal(map[string]any{ + "errorMessage": errorMessage, + }) + case "final_answer": + if finalContent, _ := payloadMap["content"].(string); strings.TrimSpace(finalContent) != "" { + assistantContent = finalContent + } + turnStatus = agentTurnStatusCompleted + } + } + handleEventBlock := func(eventType string, data []byte) { + if eventType == "" || len(data) == 0 { + return + } + persistRunEvent(eventType, data) + } + + scanner := bufio.NewScanner(resp.Body) + scanner.Buffer(make([]byte, 0, agentSSEInitialBufferSize), agentSSEMaxBufferSize) + currentEvent := "" + var currentData bytes.Buffer + clientConnected := true + for scanner.Scan() { + line := scanner.Text() + if clientConnected { + if _, writeErr := fmt.Fprintf(c.Writer, "%s\n", line); writeErr != nil { + clientConnected = false + } else { + c.Writer.Flush() + } + } + switch { + case line == "": + handleEventBlock(currentEvent, currentData.Bytes()) + currentEvent = "" + currentData.Reset() + case len(line) > 7 && line[:7] == "event: ": + currentEvent = line[7:] + case len(line) > 6 && line[:6] == "data: ": + currentData.WriteString(line[6:]) + } + } + if currentEvent != "" && currentData.Len() > 0 { + handleEventBlock(currentEvent, currentData.Bytes()) + } + if err := scanner.Err(); err != nil { + turnStatus = agentTurnStatusFailed + turnMetadata, _ = json.Marshal(map[string]any{ + "errorMessage": err.Error(), + }) + } + if turnStatus == agentTurnStatusRunning { + turnStatus = agentTurnStatusFailed + turnMetadata, _ = json.Marshal(map[string]any{ + "errorMessage": "Agent 未返回最终答复，执行可能已中断。", + }) + } + + var finalMessageID *uint + if persistAssistant && strings.TrimSpace(assistantContent) != "" { + assistantMsg := &model.AgentMessage{ + SessionID: sessionID, + Role: agentMessageRoleAssistant, + Content: assistantContent, + CreatedAt: time.Now(), + } + if err := mgr.agentService.SaveMessage(context.Background(), assistantMsg); err == nil { + finalMessageID = &assistantMsg.ID + } + } + switch turnStatus { + case agentTurnStatusFailed: + _ = mgr.agentService.UpdateTurnStatus(context.Background(), turnID, agentTurnStatusFailed, finalMessageID, turnMetadata) + case agentToolStatusAwaitingConfirmation: + _ = mgr.agentService.UpdateTurnStatus(context.Background(), turnID, agentToolStatusAwaitingConfirmation, finalMessageID, nil) + default: + _ = mgr.agentService.UpdateTurnStatus(context.Background(), turnID, agentTurnStatusCompleted, finalMessageID, nil) + } +} + +func parseToolCallResult(toolCall *model.AgentToolCall) (result any, fallback string) { + if toolCall == nil || len(toolCall.ToolResult) == 0 { + return nil, "" + } + + var payload any + if err := json.Unmarshal(toolCall.ToolResult, &payload); err != nil { + return nil, "" + } + + resultMap, _ := payload.(map[string]any) + if resultMap == nil { + return payload, "" + } + if errorMessage, _ := resultMap["error"].(string); strings.TrimSpace(errorMessage) != "" { + return payload, strings.TrimSpace(errorMessage) + } + if message, _ := resultMap["message"].(string); strings.TrimSpace(message) != "" { + return payload, strings.TrimSpace(message) + } + return payload, "" +} + +func (mgr *AgentMgr) buildConfirmationFinalAnswer(toolCall *model.AgentToolCall) string { + if toolCall == nil { + return "刚才的确认操作已经结束。" + } + if toolCall.ResultStatus == agentToolStatusRejected { + return fmt.Sprintf( + "已取消该操作（%s）。我不会继续执行或再次确认；如果需要其他处理，请重新告诉我。", + toolCall.ToolName, + ) + } + + result, fallback := parseToolCallResult(toolCall) + answer := strings.TrimSpace(mgr.buildToolOutcomeMessage(toolCall.ToolName, toolCall.ResultStatus, result, fallback)) + if answer == "" { + answer = "刚才的确认操作已经结束。" + } + if toolCall.ResultStatus == agentToolStatusError { + switch toolCall.ToolName { + case agentToolResubmitJob: + answer += " 你可以调整资源配置后再试，或先查看原作业详情确认失败原因。" + case agentToolCreateJupyter, agentToolCreateCustom: + answer += " 你可以修改表单参数后再试一次。" + } + } + return answer +} + +func (mgr *AgentMgr) persistSyntheticRunEvent( + turnID string, + sessionID string, + agentID string, + agentRole string, + eventType string, + eventStatus string, + content string, + metadata map[string]any, +) { + if turnID == "" || sessionID == "" || eventType == "" { + return + } + metadataBytes, _ := json.Marshal(metadata) + timestamp := time.Now() + _, _ = mgr.agentService.CreateRunEvent(context.Background(), &model.AgentRunEvent{ + TurnID: turnID, + SessionID: sessionID, + AgentID: agentID, + AgentRole: agentRole, + EventType: eventType, + EventStatus: eventStatus, + Content: content, + Metadata: datatypes.JSON(metadataBytes), + StartedAt: ×tamp, + EndedAt: ×tamp, + CreatedAt: timestamp, + }) +} + +func writeSyntheticSSEEvent(c *gin.Context, eventType string, payload map[string]any) error { + data, err := json.Marshal(payload) + if err != nil { + return err + } + if _, err := fmt.Fprintf(c.Writer, "event: %s\ndata: %s\n\n", eventType, data); err != nil { + return err + } + if flusher, ok := c.Writer.(http.Flusher); ok { + flusher.Flush() + } + return nil +} + +func (mgr *AgentMgr) streamConfirmationOutcome( + c *gin.Context, + sessionID string, + turnID string, + toolCall *model.AgentToolCall, +) { + finalAnswer := mgr.buildConfirmationFinalAnswer(toolCall) + agentID := "coordinator-1" + agentRole := agentRoleCoordinator + + c.Header("Content-Type", "text/event-stream") + c.Header("Cache-Control", "no-cache") + c.Header("Connection", "keep-alive") + c.Header("X-Agent-Session-ID", sessionID) + c.Header("X-Agent-Turn-ID", turnID) + c.Header("X-Agent-Orchestration-Mode", "confirmation_result") + c.Status(http.StatusOK) + + startPayload := map[string]any{ + "turnId": turnID, + "sessionId": sessionID, + "agentId": agentID, + "agentRole": agentRole, + "status": agentTurnStatusCompleted, + "summary": "确认操作已完成，正在整理结果", + "source": "confirmation_resume", + "toolName": toolCall.ToolName, + "toolStatus": toolCall.ResultStatus, + } + finalPayload := map[string]any{ + "turnId": turnID, + "sessionId": sessionID, + "agentId": agentID, + "agentRole": agentRole, + "content": finalAnswer, + "source": "confirmation_resume", + } + + mgr.persistSyntheticRunEvent( + turnID, + sessionID, + agentID, + agentRole, + "agent_run_started", + agentTurnStatusCompleted, + "确认操作已完成，正在整理结果", + startPayload, + ) + mgr.persistSyntheticRunEvent( + turnID, + sessionID, + agentID, + agentRole, + "final_answer", + agentTurnStatusCompleted, + finalAnswer, + finalPayload, + ) + _ = mgr.agentService.UpdateTurnStatus(context.Background(), turnID, agentTurnStatusCompleted, nil, nil) + + _ = writeSyntheticSSEEvent(c, "agent_run_started", startPayload) + _ = writeSyntheticSSEEvent(c, "final_answer", finalPayload) + _ = writeSyntheticSSEEvent(c, "done", map[string]any{}) +} + +// proxySSEStream is a utility that copies an SSE response from an upstream URL to the Gin client. +// It is kept for future use when direct proxying is preferred over buffered streaming. +// +//nolint:unused // reserved for future direct-proxy SSE support +func proxySSEStream(c *gin.Context, httpClient *http.Client, upstreamURL string, body io.Reader) error { + req, err := http.NewRequestWithContext(c.Request.Context(), http.MethodPost, upstreamURL, body) + if err != nil { + return bizerr.Internal.ServiceError.Wrap(err, "failed to create upstream request") + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "text/event-stream") + + resp, err := httpClient.Do(req) + if err != nil { + return bizerr.BadGateway.BadGateway.Wrap(err, "upstream request failed") + } + defer resp.Body.Close() + + c.Header("Content-Type", "text/event-stream") + c.Header("Cache-Control", "no-cache") + c.Header("Connection", "keep-alive") + + scanner := bufio.NewScanner(resp.Body) + for scanner.Scan() { + line := scanner.Text() + if _, writeErr := fmt.Fprintf(c.Writer, "%s\n", line); writeErr != nil { + return writeErr + } + c.Writer.Flush() + } + return scanner.Err() +} diff --git a/backend/internal/handler/agent/tool_outcome.go b/backend/internal/handler/agent/tool_outcome.go new file mode 100644 index 000000000..c819e1752 --- /dev/null +++ b/backend/internal/handler/agent/tool_outcome.go @@ -0,0 +1,50 @@ +package agent + +import ( + "encoding/json" + "fmt" + "strings" +) + +func (mgr *AgentMgr) buildToolOutcomeMessage(toolName, status string, result any, fallback string) string { + if strings.TrimSpace(fallback) != "" && status != agentToolStatusSuccess { + return fallback + } + + target := extractToolOutcomeTarget(result) + switch status { + case agentToolStatusSuccess: + if target != "" { + return fmt.Sprintf("%s 已完成，目标：%s。", toolName, target) + } + return fmt.Sprintf("%s 已完成。", toolName) + case agentToolStatusRejected: + return fmt.Sprintf("已取消 %s。", toolName) + case agentTurnStatusCancelled: + return fmt.Sprintf("%s 已取消。", toolName) + default: + if strings.TrimSpace(fallback) != "" { + return fallback + } + return fmt.Sprintf("%s 执行失败。", toolName) + } +} + +func extractToolOutcomeTarget(result any) string { + resultMap, ok := result.(map[string]any) + if !ok { + resultBytes, err := json.Marshal(result) + if err != nil { + return "" + } + if err := json.Unmarshal(resultBytes, &resultMap); err != nil { + return "" + } + } + for _, key := range []string{"job_name", "jobName", "name", "node", "node_name", "image"} { + if value, ok := resultMap[key].(string); ok && strings.TrimSpace(value) != "" { + return strings.TrimSpace(value) + } + } + return "" +} diff --git a/backend/internal/handler/agent/tools_dispatch.go b/backend/internal/handler/agent/tools_dispatch.go new file mode 100644 index 000000000..f42d0eb4a --- /dev/null +++ b/backend/internal/handler/agent/tools_dispatch.go @@ -0,0 +1,505 @@ +package agent + +import ( + "encoding/json" + "fmt" + "net/url" + "strconv" + "strings" + "time" + + "github.com/gin-gonic/gin" + "gorm.io/datatypes" + + "github.com/raids-lab/crater/dao/model" + "github.com/raids-lab/crater/internal/bizerr" + "github.com/raids-lab/crater/internal/resputil" + "github.com/raids-lab/crater/internal/util" +) + +func agentDispatchErrorf(format string, args ...any) error { + return bizerr.BadRequest.ParameterError.New(fmt.Sprintf(strings.ReplaceAll(format, "%w", "%v"), args...)) +} + +func isAgentReadOnlyTool(toolName string) bool { + switch toolName { + case agentToolGetJobDetail, + agentToolGetJobEvents, + agentToolGetJobLogs, + agentToolDiagnoseJob, + agentToolGetDiagnosticCtx, + agentToolSearchSimilarFail, + agentToolQueryJobMetrics, + agentToolAnalyzeQueue, + agentToolRealtimeCapacity, + agentToolListImages, + agentToolListGPUModels, + agentToolCheckQuota, + agentToolListUserJobs, + agentToolGetJobTemplates, + agentToolResourceRecommend: + return true + default: + return false + } +} + +func isAgentConfirmTool(toolName string) bool { + switch toolName { + case agentToolResubmitJob, agentToolStopJob, agentToolDeleteJob, + agentToolCreateJupyter, agentToolCreateWebIDE, agentToolCreateCustom, + agentToolCreatePytorch, agentToolCreateTensorflow: + return true + default: + return false + } +} + +func isAgentAutoActionTool(toolName string) bool { + _ = toolName + return false +} + +func isAgentAdminOnlyTool(_ string) bool { + return false +} + +func normalizeAgentRole(role string) string { + switch strings.TrimSpace(strings.ToLower(role)) { + case agentRoleCoordinator, "planner", "explorer", "executor", "verifier", "guide", "general", agentRoleSingleAgent: + return strings.TrimSpace(strings.ToLower(role)) + default: + return agentRoleSingleAgent + } +} + +func normalizeInternalToolRole(role string) string { + switch strings.TrimSpace(strings.ToLower(role)) { + case agentSessionSourceAdmin, "system_admin", "platform_admin": + return agentSessionSourceAdmin + default: + return "" + } +} + +func normalizeRequestedSessionSource(source string) string { + switch strings.TrimSpace(strings.ToLower(source)) { + case agentSessionSourceOpsAudit: + return agentSessionSourceOpsAudit + case agentSessionSourceSystem: + return agentSessionSourceSystem + default: + return agentSessionSourceChat + } +} + +func defaultInternalSessionTitle(source, toolName, providedTitle string) string { + title := strings.TrimSpace(providedTitle) + if title != "" { + return title + } + + prefix := "[system]" + if source == agentSessionSourceOpsAudit { + prefix = "[audit]" + } + + name := strings.TrimSpace(toolName) + if name == "" { + name = "internal-task" + } + return fmt.Sprintf("%s %s", prefix, name) +} + +func toolCallAuditSourceForSessionSource(sessionSource string) string { + _ = normalizeRequestedSessionSource(sessionSource) + return agentToolAuditSourceBackend +} + +func toolCallAuditSourceForExecution(sessionSource, executionBackend string) string { + _ = normalizeRequestedSessionSource(sessionSource) + _ = executionBackend + return agentToolAuditSourceBackend +} + +func normalizeExecutionBackend(toolName, executionBackend string) string { + normalized := strings.TrimSpace(strings.ToLower(executionBackend)) + if !isAgentConfirmTool(toolName) { + return "" + } + if normalized == "" || normalized == agentToolAuditSourceBackend { + return agentToolAuditSourceBackend + } + return agentToolAuditSourceBackend +} + +func agentSessionAllowsAdminTools(session *model.AgentSession) bool { + if session == nil { + return false + } + switch normalizeRequestedSessionSource(session.Source) { + case agentSessionSourceOpsAudit, agentSessionSourceSystem: + return true + } + + return agentSessionSurface(session) == agentSessionSourceAdmin +} + +func normalizeAgentSurface(raw string) string { + switch strings.TrimSpace(strings.ToLower(raw)) { + case agentSessionSourceAdmin, "administrator", "management": + return agentSessionSourceAdmin + default: + return agentSessionSourceUser + } +} + +func agentPageContextSurface(page map[string]any) string { + if page == nil { + return agentSessionSourceUser + } + if raw, _ := page["surface"].(string); normalizeAgentSurface(raw) == agentSessionSourceAdmin { + return agentSessionSourceAdmin + } + for _, key := range []string{"route", "url"} { + raw, _ := page[key].(string) + if isAdminAgentRoute(raw) { + return agentSessionSourceAdmin + } + } + return agentSessionSourceUser +} + +func agentPageScopeForToken(token util.JWTMessage, page map[string]any) string { + if token.RolePlatform != model.RoleAdmin { + return agentSessionSourceUser + } + return agentPageContextSurface(page) +} + +func agentSessionSurface(session *model.AgentSession) string { + if session == nil { + return agentSessionSourceUser + } + switch normalizeRequestedSessionSource(session.Source) { + case agentSessionSourceOpsAudit, agentSessionSourceSystem: + return agentSessionSourceAdmin + } + return agentPageContextSurface(normalizePageContext(json.RawMessage(session.PageContext))) +} + +func agentSessionMatchesSurface(session *model.AgentSession, surface string) bool { + return agentSessionSurface(session) == normalizeAgentSurface(surface) +} + +func isAdminAgentRoute(raw string) bool { + route := strings.TrimSpace(raw) + if route == "" { + return false + } + if parsed, err := url.Parse(route); err == nil && parsed.Path != "" { + route = parsed.Path + } + route = strings.TrimSpace(strings.ToLower(route)) + return route == "/admin" || strings.HasPrefix(route, "/admin/") +} + +func effectiveAgentSessionToken(session *model.AgentSession, token util.JWTMessage) util.JWTMessage { + if token.RolePlatform == model.RoleAdmin && !agentSessionAllowsAdminTools(session) { + token.RolePlatform = model.RoleUser + } + return token +} + +func authorizeAgentToolForSession(session *model.AgentSession, token util.JWTMessage, toolName string) error { + if !isAgentAdminOnlyTool(toolName) { + return nil + } + if token.RolePlatform != model.RoleAdmin { + return agentDispatchErrorf("你当前没有管理员权限，不能执行该运维操作；如确需处理，请联系平台管理员或切换到管理员页面后再操作") + } + if !agentSessionAllowsAdminTools(session) { + return agentDispatchErrorf("该运维操作只能在管理员页面执行；用户端会话不能发起节点、Pod 或集群级写操作") + } + return nil +} + +func (mgr *AgentMgr) ensureInternalAuditSession(c *gin.Context, req *ExecuteToolRequest) error { + if req.InternalContext == nil { + return nil + } + + source := normalizeRequestedSessionSource(req.SessionSource) + if source == agentSessionSourceChat { + source = agentSessionSourceSystem + } + + _, _, err := mgr.agentService.GetOrCreateSessionWithSource( + c.Request.Context(), + req.SessionID, + 0, + 0, + defaultInternalSessionTitle(source, req.ToolName, req.SessionTitle), + nil, + source, + ) + return err +} + +func (mgr *AgentMgr) resolveToolExecutionToken(c *gin.Context, req *ExecuteToolRequest) (util.JWTMessage, error) { + if req.InternalContext != nil { + if normalizeInternalToolRole(req.InternalContext.Role) != agentSessionSourceAdmin { + return util.JWTMessage{}, agentDispatchErrorf("unsupported internal tool role") + } + username := strings.TrimSpace(req.InternalContext.Username) + if username == "" { + username = "agent-pipeline" + } + accountName := strings.TrimSpace(req.InternalContext.AccountName) + if accountName == "" { + accountName = agentSessionSourceSystem + } + return util.JWTMessage{ + Username: username, + AccountName: accountName, + RoleAccount: model.RoleAdmin, + RolePlatform: model.RoleAdmin, + }, nil + } + + session, err := mgr.agentService.GetSession(c.Request.Context(), req.SessionID) + if err != nil { + return util.JWTMessage{}, agentDispatchErrorf("session not found") + } + return mgr.getSessionToken(c.Request.Context(), session) +} + +func validateAgentToolAccess(agentRole, toolName string) error { + role := normalizeAgentRole(agentRole) + + switch role { + case agentRoleCoordinator, "planner", "explorer", "verifier", "guide", "general": + if isAgentReadOnlyTool(toolName) { + return nil + } + if isAgentAutoActionTool(toolName) { + return agentDispatchErrorf("agent role '%s' cannot execute auto-action tools", role) + } + if isAgentConfirmTool(toolName) { + return agentDispatchErrorf("agent role '%s' cannot execute confirmation tools", role) + } + return agentDispatchErrorf("agent role '%s' can only execute read-only tools", role) + case "executor", agentRoleSingleAgent: + if isAgentReadOnlyTool(toolName) || isAgentConfirmTool(toolName) || isAgentAutoActionTool(toolName) { + return nil + } + return agentDispatchErrorf("tool '%s' is not supported", toolName) + default: + return agentDispatchErrorf("agent role '%s' is not allowed to execute tools", role) + } +} + +// ExecuteTool godoc +// @Summary Execute a named tool (called by the Python Agent service) +// @Description Routes tool_name to the appropriate internal handler. Write tools return confirmation_required. +// @Tags agent +// @Accept json +// @Produce json +// @Param request body ExecuteToolRequest true "Tool execution request" +// @Success 200 {object} AgentToolResponse +// @Router /api/agent/tools/execute [post] +// +//nolint:gocyclo // Tool routing dispatches many named tools in one function. +func (mgr *AgentMgr) ExecuteTool(c *gin.Context) { + if !mgr.isInternalToolRequestAuthorized(c) { + resputil.HandleError(c, bizerr.Auth.TokenInvalid.New("invalid internal agent token")) + return + } + + var req ExecuteToolRequest + if err := c.ShouldBindJSON(&req); err != nil { + agentBadRequest(c, err.Error()) + return + } + req.AgentRole = normalizeAgentRole(req.AgentRole) + if accessErr := validateAgentToolAccess(req.AgentRole, req.ToolName); accessErr != nil { + agentForbidden(c, accessErr.Error()) + return + } + if err := mgr.ensureInternalAuditSession(c, &req); err != nil { + agentInternalError(c, fmt.Sprintf("failed to create internal audit session: %v", err)) + return + } + + sessionToken, tokenErr := mgr.resolveToolExecutionToken(c, &req) + if tokenErr != nil { + agentInternalError(c, fmt.Sprintf("failed to resolve tool actor: %v", tokenErr)) + return + } + if req.InternalContext == nil { + session, sessionErr := mgr.agentService.GetSession(c.Request.Context(), req.SessionID) + if sessionErr != nil { + agentForbidden(c, "session not found") + return + } + if authErr := authorizeAgentToolForSession(session, sessionToken, req.ToolName); authErr != nil { + agentForbidden(c, authErr.Error()) + return + } + } else if isAgentAdminOnlyTool(req.ToolName) && sessionToken.RolePlatform != model.RoleAdmin { + agentForbidden(c, "你当前没有管理员权限，不能执行该运维操作；如确需处理，请联系平台管理员或切换到管理员页面后再操作") + return + } + + if isAgentConfirmTool(req.ToolName) { + if req.InternalContext == nil { + if preflightErr := mgr.validateOwnedJobMutationBeforeConfirmation(c, sessionToken, req.ToolName, req.ToolArgs); preflightErr != nil { + agentForbidden(c, preflightErr.Error()) + return + } + } + start := time.Now() + executionBackend := normalizeExecutionBackend(req.ToolName, req.ExecutionBackend) + confirmation := mgr.buildToolConfirmation(sessionToken, req.ToolName, req.ToolArgs) + pendingResult, _ := json.Marshal(map[string]any{ + "description": confirmation.Description, + "riskLevel": confirmation.RiskLevel, + "permissionExplanation": confirmation.PermissionExplanation, + "riskExplanation": confirmation.RiskExplanation, + "affectedResources": confirmation.AffectedResources, + "interaction": confirmation.Interaction, + "form": confirmation.Form, + "execution_backend": executionBackend, + }) + toolCallRecord := &model.AgentToolCall{ + SessionID: req.SessionID, + TurnID: req.TurnID, + ToolCallID: req.ToolCallID, + AgentID: req.AgentID, + AgentRole: req.AgentRole, + Source: toolCallAuditSourceForExecution(req.SessionSource, executionBackend), + ToolName: req.ToolName, + ToolArgs: datatypes.JSON(req.ToolArgs), + ToolResult: pendingResult, + ResultStatus: agentToolStatusAwaitConfirm, + ExecutionBackend: executionBackend, + CreatedAt: time.Now(), + } + toolCall, createErr := mgr.agentService.CreateToolCall(c.Request.Context(), toolCallRecord) + if createErr != nil { + agentInternalError(c, fmt.Sprintf("failed to create pending tool call: %v", createErr)) + return + } + resputil.Success(c, AgentToolResponse{ + ToolCallID: req.ToolCallID, + Status: agentToolStatusConfirmationRequired, + Confirmation: &AgentToolConfirmation{ + ConfirmID: strconv.FormatUint(uint64(toolCall.ID), 10), + ToolName: req.ToolName, + Description: confirmation.Description, + RiskLevel: confirmation.RiskLevel, + PermissionExplanation: confirmation.PermissionExplanation, + RiskExplanation: confirmation.RiskExplanation, + AffectedResources: confirmation.AffectedResources, + Interaction: confirmation.Interaction, + Form: confirmation.Form, + }, + LatencyMs: int(time.Since(start).Milliseconds()), + }) + return + } + + start := time.Now() + + result, execErr := mgr.executeReadTool(c, sessionToken, &req) + latencyMs := int(time.Since(start).Milliseconds()) + + status := agentToolStatusSuccess + var resultBytes json.RawMessage + var errMsg string + + if execErr != nil { + status = agentToolStatusError + errMsg = execErr.Error() + errJSON, _ := json.Marshal(map[string]string{"error": errMsg}) + resultBytes = errJSON + } else { + resultBytes, _ = json.Marshal(result) + } + + mgr.agentService.LogToolCallAsync( + req.SessionID, req.ToolName, + req.ToolArgs, resultBytes, + status, latencyMs, req.TurnID, req.ToolCallID, req.AgentID, req.AgentRole, + toolCallAuditSourceForSessionSource(req.SessionSource), + ) + + resputil.Success(c, AgentToolResponse{ + ToolCallID: req.ToolCallID, + Status: status, + Result: resultBytes, + Message: errMsg, + LatencyMs: latencyMs, + }) +} + +//nolint:gocyclo // Tool dispatch intentionally maps all read-only tools in one place. +func (mgr *AgentMgr) executeReadTool(c *gin.Context, token util.JWTMessage, req *ExecuteToolRequest) (any, error) { + switch req.ToolName { + case agentToolGetJobDetail: + return mgr.toolGetJobDetail(c, token, req.ToolArgs) + case agentToolGetJobEvents: + return mgr.toolGetJobEvents(c, token, req.ToolArgs) + case agentToolGetJobLogs: + return mgr.toolGetJobLogs(c, token, req.ToolArgs) + case agentToolDiagnoseJob: + return mgr.toolDiagnoseJob(c, token, req.ToolArgs) + case agentToolGetDiagnosticCtx: + return mgr.toolGetDiagnosticContext(c, token, req.ToolArgs) + case agentToolSearchSimilarFail: + return mgr.toolSearchSimilarFailures(c, token, req.ToolArgs) + case agentToolQueryJobMetrics: + return mgr.toolQueryJobMetrics(c, token, req.ToolArgs) + case agentToolAnalyzeQueue: + return mgr.toolAnalyzeQueueStatus(c, token, req.ToolArgs) + case agentToolRealtimeCapacity: + return mgr.toolGetRealtimeCapacity(c, token, req.ToolArgs) + case agentToolListImages: + return mgr.toolListAvailableImages(c, token, req.ToolArgs) + case agentToolListGPUModels: + return mgr.toolListAvailableGPUModels(c, token, req.ToolArgs) + case agentToolCheckQuota: + return mgr.toolCheckQuota(c, token, req.ToolArgs) + case agentToolListUserJobs: + return mgr.toolListUserJobs(c, token, req.ToolArgs) + case agentToolGetJobTemplates: + return mgr.toolGetJobTemplates(c, token, req.ToolArgs) + case agentToolResourceRecommend: + return mgr.toolGetResourceRecommendation(c, token, req.ToolArgs) + default: + return nil, agentDispatchErrorf("tool '%s' is not yet implemented", req.ToolName) + } +} + +func (mgr *AgentMgr) executeWriteTool(c *gin.Context, token util.JWTMessage, toolName string, rawArgs json.RawMessage) (any, error) { + switch toolName { + case agentToolDeleteJob: + return mgr.toolDeleteJob(c, token, rawArgs) + case agentToolStopJob: + return mgr.toolStopJob(c, token, rawArgs) + case agentToolResubmitJob: + return mgr.toolResubmitJob(c, token, rawArgs) + case agentToolCreateJupyter: + return mgr.toolCreateJupyterJob(c, token, rawArgs) + case agentToolCreateWebIDE: + return mgr.toolCreateWebIDEJob(c, token, rawArgs) + case agentToolCreateCustom: + return mgr.toolCreateCustomJob(c, token, rawArgs) + case agentToolCreatePytorch: + return mgr.toolCreatePytorchJob(c, token, rawArgs) + case agentToolCreateTensorflow: + return mgr.toolCreateTensorflowJob(c, token, rawArgs) + default: + return nil, agentDispatchErrorf("write tool '%s' is not supported", toolName) + } +} diff --git a/backend/internal/handler/agent/tools_gpu.go b/backend/internal/handler/agent/tools_gpu.go new file mode 100644 index 000000000..d0a12e7aa --- /dev/null +++ b/backend/internal/handler/agent/tools_gpu.go @@ -0,0 +1,33 @@ +package agent + +import ( + "fmt" + "strings" + + v1 "k8s.io/api/core/v1" +) + +func normalizeGPUModelName(input string) string { + model := strings.TrimSpace(strings.ToLower(input)) + model = strings.ReplaceAll(model, " ", "-") + return model +} + +func normalizeGPUResourceName(current v1.ResourceName, gpuModel string) v1.ResourceName { + model := normalizeGPUModelName(gpuModel) + if model == "" { + return current + } + if strings.Contains(model, "/") { + return v1.ResourceName(model) + } + + vendor := "nvidia.com" + if current != "" { + parts := strings.SplitN(string(current), "/", 2) + if len(parts) == 2 && parts[0] != "" { + vendor = parts[0] + } + } + return v1.ResourceName(fmt.Sprintf("%s/%s", vendor, model)) +} diff --git a/backend/internal/handler/agent/tools_readonly.go b/backend/internal/handler/agent/tools_readonly.go new file mode 100644 index 000000000..11a273797 --- /dev/null +++ b/backend/internal/handler/agent/tools_readonly.go @@ -0,0 +1,358 @@ +package agent + +import ( + "encoding/json" + "fmt" + "strings" + "time" + + "github.com/gin-gonic/gin" + "volcano.sh/apis/pkg/apis/batch/v1alpha1" + + "github.com/raids-lab/crater/dao/model" + "github.com/raids-lab/crater/dao/query" + basehandler "github.com/raids-lab/crater/internal/handler" + "github.com/raids-lab/crater/internal/util" +) + +const ( + agentDefaultJobLogTailLines = 200 + agentMaxJobLogTailLines = 2000 + agentDefaultDiagnosticTailLines = 120 + agentMaxDiagnosticTailLines = 1000 + agentDefaultSimilarFailureLimit = 5 + agentMaxSimilarFailureLimit = 20 + agentDefaultSimilarFailureDays = 30 + agentMaxSimilarFailureDays = 180 + agentSimilarFailureQueryMultiplier = 4 + agentDefaultListImagesLimit = 30 + agentMaxListImagesLimit = 100 + agentDefaultListJobsLimit = 20 + agentMaxListJobsLimit = 100 +) + +func agentReadonlyErrorf(format string, args ...any) error { + return agentDispatchErrorf(format, args...) +} + +func (mgr *AgentMgr) requireJobReader() (basehandler.JobInsightReader, error) { + if mgr.jobReader == nil { + return nil, agentReadonlyErrorf("job reader is not available") + } + return mgr.jobReader, nil +} + +func (mgr *AgentMgr) findScopedJob(ctx *gin.Context, token util.JWTMessage, jobName string) (*model.Job, error) { + reader, err := mgr.requireJobReader() + if err != nil { + return nil, err + } + return reader.FindScopedJob(ctx.Request.Context(), token, strings.TrimSpace(jobName)) +} + +func getJobNameArg(rawArgs json.RawMessage) (string, error) { + args := parseToolArgsMap(rawArgs) + jobName := getToolArgString(args, "job_name", "") + if jobName == "" { + jobName = getToolArgString(args, "jobName", "") + } + if jobName == "" { + return "", agentReadonlyErrorf("job_name is required") + } + return jobName, nil +} + +func (mgr *AgentMgr) toolGetJobDetail(c *gin.Context, token util.JWTMessage, rawArgs json.RawMessage) (any, error) { + jobName, err := getJobNameArg(rawArgs) + if err != nil { + return nil, err + } + reader, err := mgr.requireJobReader() + if err != nil { + return nil, err + } + job, err := reader.FindScopedJob(c.Request.Context(), token, jobName) + if err != nil { + return nil, err + } + return reader.BuildJobDetail(job), nil +} + +func (mgr *AgentMgr) toolGetJobEvents(c *gin.Context, token util.JWTMessage, rawArgs json.RawMessage) (any, error) { + jobName, err := getJobNameArg(rawArgs) + if err != nil { + return nil, err + } + reader, err := mgr.requireJobReader() + if err != nil { + return nil, err + } + return reader.GetJobEvents(c.Request.Context(), token, jobName) +} + +func (mgr *AgentMgr) toolGetJobLogs(c *gin.Context, token util.JWTMessage, rawArgs json.RawMessage) (any, error) { + args := parseToolArgsMap(rawArgs) + jobName, err := getJobNameArg(rawArgs) + if err != nil { + return nil, err + } + tailLines := int64(getToolArgInt(args, "tail", agentDefaultJobLogTailLines)) + if tailLines <= 0 || tailLines > agentMaxJobLogTailLines { + tailLines = agentDefaultJobLogTailLines + } + reader, err := mgr.requireJobReader() + if err != nil { + return nil, err + } + return reader.GetJobLog( + c.Request.Context(), + token, + jobName, + tailLines, + getToolArgString(args, "keyword", ""), + ) +} + +func (mgr *AgentMgr) toolDiagnoseJob(c *gin.Context, token util.JWTMessage, rawArgs json.RawMessage) (any, error) { + jobName, err := getJobNameArg(rawArgs) + if err != nil { + return nil, err + } + job, err := mgr.findScopedJob(c, token, jobName) + if err != nil { + return nil, err + } + return basehandler.PerformDiagnosis(job), nil +} + +func (mgr *AgentMgr) toolGetDiagnosticContext(c *gin.Context, token util.JWTMessage, rawArgs json.RawMessage) (any, error) { + args := parseToolArgsMap(rawArgs) + jobName, err := getJobNameArg(rawArgs) + if err != nil { + return nil, err + } + reader, err := mgr.requireJobReader() + if err != nil { + return nil, err + } + tailLines := int64(getToolArgInt(args, "tail", agentDefaultDiagnosticTailLines)) + if tailLines <= 0 || tailLines > agentMaxDiagnosticTailLines { + tailLines = agentDefaultDiagnosticTailLines + } + return reader.GetDiagnosticContext( + c.Request.Context(), + token, + jobName, + getToolArgBool(args, "include_log", true), + tailLines, + ) +} + +func (mgr *AgentMgr) toolQueryJobMetrics(c *gin.Context, token util.JWTMessage, rawArgs json.RawMessage) (any, error) { + jobName, err := getJobNameArg(rawArgs) + if err != nil { + return nil, err + } + job, err := mgr.findScopedJob(c, token, jobName) + if err != nil { + return nil, err + } + if job.ProfileData == nil { + return map[string]any{"job_name": job.JobName, "metrics": map[string]any{}}, nil + } + return map[string]any{"job_name": job.JobName, "metrics": job.ProfileData.Data()}, nil +} + +func (mgr *AgentMgr) toolSearchSimilarFailures(c *gin.Context, token util.JWTMessage, rawArgs json.RawMessage) (any, error) { + args := parseToolArgsMap(rawArgs) + jobName, err := getJobNameArg(rawArgs) + if err != nil { + return nil, err + } + target, err := mgr.findScopedJob(c, token, jobName) + if err != nil { + return nil, err + } + category := basehandler.CategorizeFailure(target).TypeName + limit := getToolArgInt(args, "limit", agentDefaultSimilarFailureLimit) + if limit <= 0 || limit > agentMaxSimilarFailureLimit { + limit = agentDefaultSimilarFailureLimit + } + days := getToolArgInt(args, "days", agentDefaultSimilarFailureDays) + if days <= 0 || days > agentMaxSimilarFailureDays { + days = agentDefaultSimilarFailureDays + } + + j := query.Job + dbQuery := j.WithContext(c.Request.Context()). + Where(j.Status.Eq(string(v1alpha1.Failed))). + Where(j.CreationTimestamp.Gte(time.Now().AddDate(0, 0, -days))) + if token.RolePlatform != model.RoleAdmin { + dbQuery = dbQuery.Where(j.UserID.Eq(token.UserID), j.AccountID.Eq(token.AccountID)) + } + jobs, err := dbQuery.Order(j.CreationTimestamp.Desc()).Limit(limit * agentSimilarFailureQueryMultiplier).Find() + if err != nil { + return nil, err + } + items := make([]map[string]any, 0, limit) + for _, job := range jobs { + if job.JobName == target.JobName { + continue + } + if basehandler.CategorizeFailure(job).TypeName != category { + continue + } + items = append(items, map[string]any{ + "job_name": job.JobName, + "name": job.Name, + "status": job.Status, + "createdAt": job.CreationTimestamp, + "category": category, + }) + if len(items) >= limit { + break + } + } + return map[string]any{"category": category, "items": items}, nil +} + +func (mgr *AgentMgr) toolAnalyzeQueueStatus(c *gin.Context, token util.JWTMessage, rawArgs json.RawMessage) (any, error) { + jobName, err := getJobNameArg(rawArgs) + if err != nil { + return nil, err + } + job, err := mgr.findScopedJob(c, token, jobName) + if err != nil { + return nil, err + } + result := map[string]any{ + "job_name": job.JobName, + "status": job.Status, + "pending": job.Status == v1alpha1.Pending, + } + if job.Events != nil { + result["events"] = job.Events.Data() + } + if job.ScheduleData != nil { + result["schedule_data"] = job.ScheduleData.Data() + } + return result, nil +} + +func (mgr *AgentMgr) toolGetRealtimeCapacity(c *gin.Context, _ util.JWTMessage, _ json.RawMessage) (any, error) { + if mgr.nodeClient == nil { + return map[string]any{"nodes": []any{}}, nil + } + nodes, err := mgr.nodeClient.ListNodes(c.Request.Context()) + if err != nil { + return nil, err + } + return map[string]any{"nodes": nodes}, nil +} + +func (mgr *AgentMgr) toolListAvailableImages(c *gin.Context, token util.JWTMessage, rawArgs json.RawMessage) (any, error) { + if mgr.imageReader == nil { + return nil, agentReadonlyErrorf("image reader is not available") + } + args := parseToolArgsMap(rawArgs) + limit := getToolArgInt(args, "limit", agentDefaultListImagesLimit) + if limit <= 0 || limit > agentMaxListImagesLimit { + limit = agentDefaultListImagesLimit + } + records, err := mgr.imageReader.ListAccessibleImages(c.Request.Context(), token) + if err != nil { + return nil, err + } + items := make([]map[string]any, 0, limit) + for _, record := range records { + if record.Image == nil { + continue + } + imageName := record.Image.ImageLink + if record.Image.ImagePackName != nil && strings.TrimSpace(*record.Image.ImagePackName) != "" { + imageName = *record.Image.ImagePackName + } + items = append(items, map[string]any{ + "id": record.Image.ID, + "name": imageName, + "image_link": record.Image.ImageLink, + "description": record.Image.Description, + "share_status": record.ShareStatus, + }) + if len(items) >= limit { + break + } + } + return map[string]any{"images": items, "count": len(items)}, nil +} + +func (mgr *AgentMgr) toolListAvailableGPUModels(_ *gin.Context, _ util.JWTMessage, _ json.RawMessage) (any, error) { + return map[string]any{"models": []string{"V100", "A100", "H100", "L40S", "RTX4090"}}, nil +} + +func (mgr *AgentMgr) toolCheckQuota(_ *gin.Context, token util.JWTMessage, _ json.RawMessage) (any, error) { + return map[string]any{ + "account_id": token.AccountID, + "user_id": token.UserID, + "message": "quota detail is available from the platform quota APIs", + }, nil +} + +func (mgr *AgentMgr) toolListUserJobs(c *gin.Context, token util.JWTMessage, rawArgs json.RawMessage) (any, error) { + args := parseToolArgsMap(rawArgs) + limit := getToolArgInt(args, "limit", agentDefaultListJobsLimit) + if limit <= 0 || limit > agentMaxListJobsLimit { + limit = agentDefaultListJobsLimit + } + j := query.Job + dbQuery := j.WithContext(c.Request.Context()).Order(j.CreationTimestamp.Desc()).Limit(limit) + if token.RolePlatform != model.RoleAdmin { + dbQuery = dbQuery.Where(j.UserID.Eq(token.UserID), j.AccountID.Eq(token.AccountID)) + } + jobs, err := dbQuery.Find() + if err != nil { + return nil, err + } + items := make([]map[string]any, 0, len(jobs)) + for _, job := range jobs { + items = append(items, map[string]any{ + "job_name": job.JobName, + "name": job.Name, + "type": job.JobType, + "status": job.Status, + "createdAt": job.CreationTimestamp, + }) + } + return map[string]any{"jobs": items, "count": len(items)}, nil +} + +func (mgr *AgentMgr) toolGetJobTemplates(_ *gin.Context, _ util.JWTMessage, _ json.RawMessage) (any, error) { + return map[string]any{ + "templates": []map[string]any{ + {"type": "jupyter", "tool": agentToolCreateJupyter}, + {"type": "webide", "tool": agentToolCreateWebIDE}, + {"type": "custom", "tool": agentToolCreateCustom}, + {"type": "pytorch", "tool": agentToolCreatePytorch}, + {"type": "tensorflow", "tool": agentToolCreateTensorflow}, + }, + }, nil +} + +func (mgr *AgentMgr) toolGetResourceRecommendation(_ *gin.Context, _ util.JWTMessage, rawArgs json.RawMessage) (any, error) { + args := parseToolArgsMap(rawArgs) + taskType := strings.ToLower(getToolArgString(args, "task_type", "")) + cpu := "4" + memory := "16Gi" + gpuCount := 0 + if strings.Contains(taskType, "train") || strings.Contains(taskType, "训练") { + cpu = "8" + memory = "32Gi" + gpuCount = 1 + } + return map[string]any{ + "cpu": cpu, + "memory": memory, + "gpu_count": gpuCount, + "reason": fmt.Sprintf("basic recommendation for %q", taskType), + }, nil +} diff --git a/backend/internal/handler/agent/tools_write.go b/backend/internal/handler/agent/tools_write.go new file mode 100644 index 000000000..fc498730b --- /dev/null +++ b/backend/internal/handler/agent/tools_write.go @@ -0,0 +1,643 @@ +package agent + +import ( + "encoding/json" + "fmt" + "strconv" + "strings" + + "github.com/gin-gonic/gin" + + "github.com/raids-lab/crater/internal/bizerr" + "github.com/raids-lab/crater/internal/util" +) + +const ( + agentForwardTypeIngress = 1 + agentForwardTypeNodePort = 2 +) + +func agentWriteErrorf(format string, args ...any) error { + return bizerr.BadRequest.ParameterError.New(fmt.Sprintf(strings.ReplaceAll(format, "%w", "%v"), args...)) +} + +func normalizeForwardTypeValue(value any) (int, error) { + switch typed := value.(type) { + case nil: + return agentForwardTypeIngress, nil + case float64: + switch int(typed) { + case agentForwardTypeIngress, agentForwardTypeNodePort: + return int(typed), nil + } + case int: + switch typed { + case agentForwardTypeIngress, agentForwardTypeNodePort: + return typed, nil + } + case string: + switch strings.TrimSpace(strings.ToLower(typed)) { + case "", "ingress": + return agentForwardTypeIngress, nil + case "nodeport", "node_port": + return agentForwardTypeNodePort, nil + case "1": + return agentForwardTypeIngress, nil + case "2": + return agentForwardTypeNodePort, nil + } + } + return 0, agentWriteErrorf("forward type must be ingress or nodeport") +} + +func parseForwardTextSpecs(raw string) ([]map[string]any, error) { + raw = strings.TrimSpace(raw) + if raw == "" { + return nil, nil + } + fields := strings.FieldsFunc(raw, func(r rune) bool { + return r == '\n' || r == '\r' || r == ',' + }) + result := make([]map[string]any, 0, len(fields)) + for _, field := range fields { + field = strings.TrimSpace(field) + if field == "" { + continue + } + parts := strings.Split(field, ":") + if len(parts) < 2 || len(parts) > 3 { + return nil, agentWriteErrorf("invalid forward spec %q, expected name:port[:ingress|nodeport]", field) + } + port, err := strconv.Atoi(strings.TrimSpace(parts[1])) + if err != nil || port <= 0 { + return nil, agentWriteErrorf("invalid forward port in %q", field) + } + forwardType := agentForwardTypeIngress + if len(parts) == 3 { + forwardType, err = normalizeForwardTypeValue(parts[2]) + if err != nil { + return nil, err + } + } + result = append(result, map[string]any{ + "type": forwardType, + "name": strings.TrimSpace(parts[0]), + "port": port, + }) + } + return result, nil +} + +func parseForwardArgs(args map[string]any) ([]map[string]any, error) { + raw, ok := args["forwards"] + if ok && raw != nil { + items, ok := raw.([]any) + if !ok { + return nil, agentWriteErrorf("forwards must be a list") + } + result := make([]map[string]any, 0, len(items)) + for _, item := range items { + entry, ok := item.(map[string]any) + if !ok { + return nil, agentWriteErrorf("forwards entries must be objects") + } + name := getToolArgString(entry, "name", "") + if name == "" { + return nil, agentWriteErrorf("forward name is required") + } + port := getToolArgInt(entry, "port", 0) + if port <= 0 { + return nil, agentWriteErrorf("forward %q requires a positive port", name) + } + forwardType, err := normalizeForwardTypeValue(entry["type"]) + if err != nil { + return nil, err + } + result = append(result, map[string]any{ + "type": forwardType, + "name": name, + "port": port, + }) + } + return result, nil + } + return parseForwardTextSpecs(getToolArgString(args, "forwards_text", "")) +} + +func lookupToolArgValue(args map[string]any, keys ...string) (any, bool) { + for _, key := range keys { + value, ok := args[key] + if !ok || value == nil { + continue + } + return value, true + } + return nil, false +} + +func getToolArgStringAny(args map[string]any, fallback string, keys ...string) string { + for _, key := range keys { + if value := getToolArgString(args, key, ""); value != "" { + return value + } + } + return fallback +} + +func getToolArgIntAny(args map[string]any, fallback int, keys ...string) int { + for _, key := range keys { + if value, ok := lookupToolArgValue(args, key); ok { + switch typed := value.(type) { + case float64: + return int(typed) + case int: + return typed + case int32: + return int(typed) + case int64: + return int(typed) + case string: + parsed, err := strconv.Atoi(strings.TrimSpace(typed)) + if err == nil { + return parsed + } + } + } + } + return fallback +} + +func parseDistributedPorts(raw any) ([]map[string]any, error) { + if raw == nil { + return nil, nil + } + + if text, ok := raw.(string); ok { + text = strings.TrimSpace(text) + if text == "" { + return nil, nil + } + if strings.HasPrefix(text, "[") { + var decoded []any + if err := json.Unmarshal([]byte(text), &decoded); err != nil { + return nil, agentWriteErrorf("ports_json must be a JSON array: %w", err) + } + raw = decoded + } else { + specs, err := parseForwardTextSpecs(text) + if err != nil { + return nil, err + } + ports := make([]map[string]any, 0, len(specs)) + for _, spec := range specs { + ports = append(ports, map[string]any{ + "name": spec["name"], + "port": spec["port"], + }) + } + return ports, nil + } + } + + items, ok := raw.([]any) + if !ok { + return nil, agentWriteErrorf("ports must be a list or text specs") + } + + ports := make([]map[string]any, 0, len(items)) + for idx, item := range items { + entry, ok := item.(map[string]any) + if !ok { + return nil, agentWriteErrorf("port #%d must be an object", idx+1) + } + name := getToolArgStringAny(entry, "", "name") + if name == "" { + return nil, agentWriteErrorf("port #%d requires name", idx+1) + } + port := getToolArgIntAny(entry, 0, "port") + if port <= 0 { + return nil, agentWriteErrorf("port %q requires a positive port number", name) + } + ports = append(ports, map[string]any{ + "name": name, + "port": port, + }) + } + return ports, nil +} + +func normalizeDistributedTask(entry map[string]any) (map[string]any, error) { + name := getToolArgStringAny(entry, "", "name") + if name == "" { + return nil, agentWriteErrorf("task name is required") + } + imageLink := getToolArgStringAny(entry, "", "image_link", "imageLink") + if imageLink == "" { + return nil, agentWriteErrorf("task %q requires image_link", name) + } + + replicas := getToolArgIntAny(entry, 1, "replicas") + if replicas <= 0 { + return nil, agentWriteErrorf("task %q requires replicas > 0", name) + } + + resourceMap := map[string]string{ + "cpu": getToolArgStringAny(entry, "4", "cpu"), + "memory": getToolArgStringAny(entry, "16Gi", "memory"), + } + if gpuCount := getToolArgIntAny(entry, 0, "gpu_count", "gpuCount"); gpuCount > 0 { + gpuResourceName := normalizeGPUResourceName("", getToolArgStringAny(entry, "gpu", "gpu_model", "gpuModel")) + resourceMap[string(gpuResourceName)] = strconv.Itoa(gpuCount) + } + + ports, err := parseDistributedPorts(func() any { + if value, ok := lookupToolArgValue(entry, "ports"); ok { + return value + } + if value, ok := lookupToolArgValue(entry, "ports_text", "ports_json"); ok { + return value + } + return nil + }()) + if err != nil { + return nil, err + } + + task := map[string]any{ + "name": name, + "replicas": replicas, + "resource": resourceMap, + "image": map[string]any{ + "imageLink": imageLink, + "archs": []string{}, + }, + "ports": ports, + } + + if command := getToolArgStringAny(entry, "", "command"); command != "" { + task["command"] = command + task["shell"] = getToolArgStringAny(entry, "bash", "shell") + } + if workingDir := getToolArgStringAny(entry, "", "working_dir", "workingDir"); workingDir != "" { + task["workingDir"] = workingDir + } + + return task, nil +} + +func parseDistributedTasks(args map[string]any) ([]map[string]any, error) { + if raw, ok := lookupToolArgValue(args, "tasks"); ok { + items, ok := raw.([]any) + if !ok { + return nil, agentWriteErrorf("tasks must be a list") + } + result := make([]map[string]any, 0, len(items)) + for idx, item := range items { + entry, ok := item.(map[string]any) + if !ok { + return nil, agentWriteErrorf("task #%d must be an object", idx+1) + } + task, err := normalizeDistributedTask(entry) + if err != nil { + return nil, err + } + result = append(result, task) + } + return result, nil + } + + rawJSON := getToolArgString(args, "tasks_json", "") + if rawJSON == "" { + return nil, nil + } + var items []any + if err := json.Unmarshal([]byte(rawJSON), &items); err != nil { + return nil, agentWriteErrorf("tasks_json must be a JSON array: %w", err) + } + result := make([]map[string]any, 0, len(items)) + for idx, item := range items { + entry, ok := item.(map[string]any) + if !ok { + return nil, agentWriteErrorf("task #%d must be an object", idx+1) + } + task, err := normalizeDistributedTask(entry) + if err != nil { + return nil, err + } + result = append(result, task) + } + return result, nil +} + +func (mgr *AgentMgr) toolDeleteJob(c *gin.Context, token util.JWTMessage, rawArgs json.RawMessage) (any, error) { + if mgr.jobSubmitter == nil { + return nil, agentWriteErrorf("job submitter is not configured") + } + var args agentJobNameArgs + if err := json.Unmarshal(rawArgs, &args); err != nil { + return nil, agentWriteErrorf("invalid args: %w", err) + } + return mgr.jobSubmitter.DeleteJob(c.Request.Context(), token, args.JobName) +} + +func (mgr *AgentMgr) toolStopJob(c *gin.Context, token util.JWTMessage, rawArgs json.RawMessage) (any, error) { + if mgr.jobSubmitter == nil { + return nil, agentWriteErrorf("job submitter is not configured") + } + var args agentJobNameArgs + if err := json.Unmarshal(rawArgs, &args); err != nil { + return nil, agentWriteErrorf("invalid args: %w", err) + } + return mgr.jobSubmitter.StopJob(c.Request.Context(), token, args.JobName) +} + +func (mgr *AgentMgr) toolResubmitJob(c *gin.Context, token util.JWTMessage, rawArgs json.RawMessage) (any, error) { + if mgr.jobSubmitter == nil { + return nil, agentWriteErrorf("job submitter is not configured") + } + return mgr.jobSubmitter.ResubmitJob(c.Request.Context(), token, rawArgs) +} + +func (mgr *AgentMgr) toolCreateJupyterJob(c *gin.Context, token util.JWTMessage, rawArgs json.RawMessage) (any, error) { + argsMap := parseToolArgsMap(rawArgs) + var args struct { + Name string `json:"name"` + ImageLink string `json:"image_link"` + CPU string `json:"cpu"` + Memory string `json:"memory"` + GPUCount *int `json:"gpu_count"` + GPUModel *string `json:"gpu_model"` + } + if err := json.Unmarshal(rawArgs, &args); err != nil { + return nil, agentWriteErrorf("invalid args: %w", err) + } + if args.Name == "" { + return nil, agentWriteErrorf("name is required") + } + if args.ImageLink == "" { + return nil, agentWriteErrorf("image_link is required") + } + if args.CPU == "" { + args.CPU = "2" + } + if args.Memory == "" { + args.Memory = "8Gi" + } + + if mgr.jobSubmitter == nil { + return nil, agentWriteErrorf("job submitter is not configured") + } + forwards, err := parseForwardArgs(argsMap) + if err != nil { + return nil, err + } + + resourceMap := map[string]string{ + "cpu": args.CPU, + "memory": args.Memory, + } + if args.GPUCount != nil && *args.GPUCount > 0 { + gpuResourceName := normalizeGPUResourceName("", "gpu") + if args.GPUModel != nil && strings.TrimSpace(*args.GPUModel) != "" { + gpuResourceName = normalizeGPUResourceName(gpuResourceName, *args.GPUModel) + } + resourceMap[string(gpuResourceName)] = strconv.Itoa(*args.GPUCount) + } + + requestBody, err := json.Marshal(map[string]any{ + "name": args.Name, + "resource": resourceMap, + "image": map[string]any{ + "imageLink": args.ImageLink, + "archs": []string{}, + }, + "forwards": forwards, + }) + if err != nil { + return nil, agentWriteErrorf("failed to marshal jupyter request: %w", err) + } + + result, err := mgr.jobSubmitter.SubmitJupyterJob(c, token, requestBody) + if err != nil { + return nil, err + } + return map[string]any{ + "status": "created", + "job": result, + }, nil +} + +func (mgr *AgentMgr) toolCreateWebIDEJob(c *gin.Context, token util.JWTMessage, rawArgs json.RawMessage) (any, error) { + argsMap := parseToolArgsMap(rawArgs) + var args struct { + Name string `json:"name"` + ImageLink string `json:"image_link"` + CPU string `json:"cpu"` + Memory string `json:"memory"` + GPUCount *int `json:"gpu_count"` + GPUModel *string `json:"gpu_model"` + } + if err := json.Unmarshal(rawArgs, &args); err != nil { + return nil, agentWriteErrorf("invalid args: %w", err) + } + if strings.TrimSpace(args.Name) == "" { + return nil, agentWriteErrorf("name is required") + } + if strings.TrimSpace(args.ImageLink) == "" { + return nil, agentWriteErrorf("image_link is required") + } + if strings.TrimSpace(args.CPU) == "" { + args.CPU = "2" + } + if strings.TrimSpace(args.Memory) == "" { + args.Memory = "8Gi" + } + if mgr.jobSubmitter == nil { + return nil, agentWriteErrorf("job submitter is not configured") + } + forwards, err := parseForwardArgs(argsMap) + if err != nil { + return nil, err + } + + resourceMap := map[string]string{ + "cpu": args.CPU, + "memory": args.Memory, + } + if args.GPUCount != nil && *args.GPUCount > 0 { + gpuResourceName := normalizeGPUResourceName("", "gpu") + if args.GPUModel != nil && strings.TrimSpace(*args.GPUModel) != "" { + gpuResourceName = normalizeGPUResourceName(gpuResourceName, *args.GPUModel) + } + resourceMap[string(gpuResourceName)] = strconv.Itoa(*args.GPUCount) + } + + requestBody, err := json.Marshal(map[string]any{ + "name": args.Name, + "resource": resourceMap, + "image": map[string]any{ + "imageLink": args.ImageLink, + "archs": []string{}, + }, + "forwards": forwards, + }) + if err != nil { + return nil, agentWriteErrorf("failed to marshal webide request: %w", err) + } + + result, err := mgr.jobSubmitter.SubmitWebIDEJob(c, token, requestBody) + if err != nil { + return nil, err + } + return map[string]any{ + "status": "created", + "job": result, + }, nil +} + +//nolint:gocyclo // Custom-job creation validates many optional form fields before delegating to vcjob submitter. +func (mgr *AgentMgr) toolCreateCustomJob(c *gin.Context, token util.JWTMessage, rawArgs json.RawMessage) (any, error) { + argsMap := parseToolArgsMap(rawArgs) + var args struct { + Name string `json:"name"` + ImageLink string `json:"image_link"` + Command string `json:"command"` + WorkingDir string `json:"working_dir"` + CPU string `json:"cpu"` + Memory string `json:"memory"` + GPUCount *int `json:"gpu_count"` + GPUModel *string `json:"gpu_model"` + Shell string `json:"shell"` + } + if err := json.Unmarshal(rawArgs, &args); err != nil { + return nil, agentWriteErrorf("invalid args: %w", err) + } + if strings.TrimSpace(args.Name) == "" { + return nil, agentWriteErrorf("name is required") + } + if strings.TrimSpace(args.ImageLink) == "" { + return nil, agentWriteErrorf("image_link is required") + } + if strings.TrimSpace(args.Command) == "" { + return nil, agentWriteErrorf("command is required") + } + if strings.TrimSpace(args.WorkingDir) == "" { + return nil, agentWriteErrorf("working_dir is required") + } + if strings.TrimSpace(args.CPU) == "" { + args.CPU = "4" + } + if strings.TrimSpace(args.Memory) == "" { + args.Memory = "16Gi" + } + if strings.TrimSpace(args.Shell) == "" { + args.Shell = "bash" + } + + resourceMap := map[string]string{ + "cpu": args.CPU, + "memory": args.Memory, + } + if args.GPUCount != nil && *args.GPUCount > 0 { + gpuResourceName := normalizeGPUResourceName("", "gpu") + if args.GPUModel != nil && strings.TrimSpace(*args.GPUModel) != "" { + gpuResourceName = normalizeGPUResourceName("", *args.GPUModel) + } + resourceMap[string(gpuResourceName)] = strconv.Itoa(*args.GPUCount) + } + + if mgr.jobSubmitter == nil { + return nil, agentWriteErrorf("job submitter is not configured") + } + forwards, err := parseForwardArgs(argsMap) + if err != nil { + return nil, err + } + + requestBody, err := json.Marshal(map[string]any{ + "name": args.Name, + "resource": resourceMap, + "workingDir": args.WorkingDir, + "command": args.Command, + "shell": args.Shell, + "image": map[string]any{ + "imageLink": args.ImageLink, + "archs": []string{}, + }, + "forwards": forwards, + }) + if err != nil { + return nil, agentWriteErrorf("failed to marshal custom job request: %w", err) + } + + result, err := mgr.jobSubmitter.SubmitTrainingJob(c, token, requestBody) + if err != nil { + return nil, err + } + + return map[string]any{ + "status": "created", + "job": result, + }, nil +} + +func (mgr *AgentMgr) toolCreatePytorchJob(c *gin.Context, token util.JWTMessage, rawArgs json.RawMessage) (any, error) { + return mgr.toolCreateDistributedJob(c, token, rawArgs, agentToolCreatePytorch) +} + +func (mgr *AgentMgr) toolCreateTensorflowJob(c *gin.Context, token util.JWTMessage, rawArgs json.RawMessage) (any, error) { + return mgr.toolCreateDistributedJob(c, token, rawArgs, agentToolCreateTensorflow) +} + +func (mgr *AgentMgr) toolCreateDistributedJob( + c *gin.Context, + token util.JWTMessage, + rawArgs json.RawMessage, + toolName string, +) (any, error) { + argsMap := parseToolArgsMap(rawArgs) + name := getToolArgString(argsMap, "name", "") + if strings.TrimSpace(name) == "" { + return nil, agentWriteErrorf("name is required") + } + tasks, err := parseDistributedTasks(argsMap) + if err != nil { + return nil, err + } + if len(tasks) == 0 { + return nil, agentWriteErrorf("tasks or tasks_json is required") + } + forwards, err := parseForwardArgs(argsMap) + if err != nil { + return nil, err + } + if mgr.jobSubmitter == nil { + return nil, agentWriteErrorf("job submitter is not configured") + } + + requestBody, err := json.Marshal(map[string]any{ + "name": name, + "tasks": tasks, + "forwards": forwards, + }) + if err != nil { + return nil, agentWriteErrorf("failed to marshal %s request: %w", toolName, err) + } + + var result any + switch toolName { + case agentToolCreatePytorch: + result, err = mgr.jobSubmitter.SubmitPytorchJob(c, token, requestBody) + case agentToolCreateTensorflow: + result, err = mgr.jobSubmitter.SubmitTensorflowJob(c, token, requestBody) + default: + return nil, agentWriteErrorf("distributed job tool %s is not supported", toolName) + } + if err != nil { + return nil, err + } + return map[string]any{ + "status": "created", + "job": result, + }, nil +} diff --git a/backend/internal/handler/agent/types.go b/backend/internal/handler/agent/types.go new file mode 100644 index 000000000..be18e1e6f --- /dev/null +++ b/backend/internal/handler/agent/types.go @@ -0,0 +1,127 @@ +package agent + +import "encoding/json" + +type AgentChatRequest struct { + Message string `json:"message" binding:"required"` + SessionID string `json:"sessionId,omitempty"` + RequestID string `json:"requestId,omitempty"` + PageContext json.RawMessage `json:"pageContext,omitempty"` + ClientContext json.RawMessage `json:"clientContext,omitempty"` + OrchestrationMode string `json:"orchestrationMode,omitempty"` +} + +type AgentAskRequest struct { + Message string `json:"message" binding:"required"` + SessionID string `json:"sessionId,omitempty"` + RequestID string `json:"requestId,omitempty"` + PageContext json.RawMessage `json:"pageContext,omitempty"` + ClientContext json.RawMessage `json:"clientContext,omitempty"` + JobName string `json:"jobName,omitempty"` +} + +type ConfirmToolRequest struct { + ConfirmID string `json:"confirmId" binding:"required"` + Confirmed bool `json:"confirmed"` + Payload json.RawMessage `json:"payload,omitempty"` +} + +type AgentResumeRequest struct { + ConfirmID string `json:"confirmId" binding:"required"` +} + +type AgentSessionPinRequest struct { + Pinned bool `json:"pinned"` +} + +type AgentSessionTitleRequest struct { + Title string `json:"title" binding:"required"` +} + +type ExecuteToolRequest struct { + ToolName string `json:"tool_name" binding:"required"` + ToolArgs json.RawMessage `json:"tool_args" binding:"required"` + SessionID string `json:"session_id" binding:"required"` + SessionSource string `json:"session_source,omitempty"` + SessionTitle string `json:"session_title,omitempty"` + ExecutionBackend string `json:"execution_backend,omitempty"` + TurnID string `json:"turn_id,omitempty"` + ToolCallID string `json:"tool_call_id,omitempty"` + AgentID string `json:"agent_id,omitempty"` + AgentRole string `json:"agent_role,omitempty"` + InternalContext *AgentInternalContext `json:"internal_context,omitempty"` +} + +type AgentInternalContext struct { + Role string `json:"role,omitempty"` + Username string `json:"username,omitempty"` + AccountName string `json:"account_name,omitempty"` +} + +type AgentConfigSummary struct { + DefaultOrchestrationMode string `json:"defaultOrchestrationMode"` + AvailableModes []string `json:"availableModes,omitempty"` + LLM *AgentLLMConfigSummary `json:"llm,omitempty"` +} + +type AgentLLMConfigSummary struct { + Source string `json:"source"` + UsingConfig bool `json:"usingConfig"` + UsingPersonal bool `json:"usingPersonal"` + Complete bool `json:"complete"` + BaseURL string `json:"baseUrl,omitempty"` + Model string `json:"model,omitempty"` + HasAPIKey bool `json:"hasApiKey"` + FallbackNote string `json:"fallbackNote,omitempty"` +} + +type AgentTurnRequest struct { + SessionID string `json:"session_id"` + TurnID string `json:"turn_id"` + Message string `json:"message"` + Context map[string]any `json:"context"` +} + +type AgentToolConfirmation struct { + ConfirmID string `json:"confirm_id"` + ToolName string `json:"tool_name"` + Description string `json:"description"` + RiskLevel string `json:"risk_level"` + PermissionExplanation string `json:"permission_explanation,omitempty"` + RiskExplanation string `json:"risk_explanation,omitempty"` + AffectedResources []string `json:"affected_resources,omitempty"` + Interaction string `json:"interaction,omitempty"` + Form *AgentToolForm `json:"form,omitempty"` +} + +type AgentToolForm struct { + Title string `json:"title,omitempty"` + Description string `json:"description,omitempty"` + SubmitLabel string `json:"submitLabel,omitempty"` + Fields []AgentToolField `json:"fields,omitempty"` +} + +type AgentToolField struct { + Key string `json:"key"` + Label string `json:"label"` + Type string `json:"type"` + Required bool `json:"required,omitempty"` + Description string `json:"description,omitempty"` + Placeholder string `json:"placeholder,omitempty"` + DefaultValue any `json:"defaultValue,omitempty"` + Options []AgentToolFieldOption `json:"options,omitempty"` +} + +type AgentToolFieldOption struct { + Value string `json:"value"` + Label string `json:"label"` +} + +type AgentToolResponse struct { + ToolCallID string `json:"tool_call_id,omitempty"` + Status string `json:"status"` + Result json.RawMessage `json:"result,omitempty"` + Message string `json:"message,omitempty"` + Confirmation *AgentToolConfirmation `json:"confirmation,omitempty"` + LatencyMs int `json:"latency_ms,omitempty"` +} diff --git a/backend/internal/handler/agent_diagnosis.go b/backend/internal/handler/agent_diagnosis.go new file mode 100644 index 000000000..9f7ca4d4b --- /dev/null +++ b/backend/internal/handler/agent_diagnosis.go @@ -0,0 +1,258 @@ +package handler + +import ( + "fmt" + "strings" + + v1 "k8s.io/api/core/v1" + batch "volcano.sh/apis/pkg/apis/batch/v1alpha1" + + "github.com/raids-lab/crater/dao/model" +) + +const ( + diagnosisConfidenceHigh = "high" + diagnosisSeverityError = "error" + diagnosisSeverityWarn = "warning" + diagnosisSeverityCrit = "critical" + + exitCodeSegmentationFault = 139 + exitCodeCommandNotFound = 127 + exitCodeGracefulTerm = 143 + maxEvidenceEventsLimit = 5 +) + +type ClassifyResult struct { + TypeName string + Sample string +} + +type DiagnosisResp struct { + JobName string `json:"jobName"` + Status string `json:"status"` + Category string `json:"category"` + Diagnosis string `json:"diagnosis"` + Solution string `json:"solution"` + Confidence string `json:"confidence"` + Severity string `json:"severity"` + Evidence struct { + ExitCode int32 `json:"exitCode,omitempty"` + ExitReason string `json:"exitReason,omitempty"` + Events []string `json:"events,omitempty"` + } `json:"evidence"` +} + +// CategorizeFailure classifies a failed job without depending on the removed diagnostics routes. +// +//nolint:gocyclo // Rule matching intentionally keeps failure classification in one place. +func CategorizeFailure(job *model.Job) ClassifyResult { + if job.TerminatedStates != nil { + terminated := job.TerminatedStates.Data() + for i := range terminated { + ts := &terminated[i] + if strings.EqualFold(ts.Reason, "OOMKilled") { + return ClassifyResult{TypeName: "OOMKilled", Sample: sampleTerminated(ts)} + } + if ts.ExitCode == exitCodeSegmentationFault { + return ClassifyResult{TypeName: "SegmentationFault", Sample: sampleTerminated(ts)} + } + if ts.ExitCode == exitCodeCommandNotFound { + return ClassifyResult{TypeName: "CommandNotFound", Sample: sampleTerminated(ts)} + } + if ts.ExitCode == exitCodeGracefulTerm { + return ClassifyResult{TypeName: "GracefulTermination", Sample: sampleTerminated(ts)} + } + } + } + if job.Events != nil { + events := job.Events.Data() + for i := range events { + ev := &events[i] + if ev.Reason == "ErrImagePull" || ev.Reason == "ImagePullBackOff" { + return ClassifyResult{TypeName: "ImagePullError", Sample: sampleEvent(ev)} + } + if ev.Reason == "FailedScheduling" { + msg := strings.ToLower(ev.Message) + switch { + case strings.Contains(msg, "insufficient"): + return ClassifyResult{TypeName: "SchedulingInsufficientResources", Sample: sampleEvent(ev)} + case strings.Contains(msg, "didn't match node selector") || strings.Contains(msg, "node(s) didn't match"): + return ClassifyResult{TypeName: "SchedulingNodeSelectorMismatch", Sample: sampleEvent(ev)} + case strings.Contains(msg, "taint"): + return ClassifyResult{TypeName: "SchedulingTaintMismatch", Sample: sampleEvent(ev)} + default: + return ClassifyResult{TypeName: "SchedulingFailed", Sample: sampleEvent(ev)} + } + } + if ev.Reason == "CrashLoopBackOff" || + (ev.Reason == "BackOff" && strings.Contains(strings.ToLower(ev.Message), "back-off restarting failed container")) { + return ClassifyResult{TypeName: "CrashLoopBackOff", Sample: sampleEvent(ev)} + } + if ev.Reason == "Evicted" { + return ClassifyResult{TypeName: "Evicted", Sample: sampleEvent(ev)} + } + if ev.Reason == "FailedMount" || strings.Contains(strings.ToLower(ev.Message), "mountvolume") { + return ClassifyResult{TypeName: "VolumeMountFailed", Sample: sampleEvent(ev)} + } + if ev.Reason == "DeadlineExceeded" { + return ClassifyResult{TypeName: "JobDeadlineExceeded", Sample: sampleEvent(ev)} + } + } + } + if job.TerminatedStates != nil { + terminated := job.TerminatedStates.Data() + for i := range terminated { + ts := &terminated[i] + if strings.EqualFold(ts.Reason, "Error") && ts.ExitCode != 0 { + return ClassifyResult{TypeName: "ContainerError", Sample: sampleTerminated(ts)} + } + } + } + switch job.Status { + case batch.Aborted, batch.Terminated: + return ClassifyResult{TypeName: "JobAbortedOrTerminated"} + } + return ClassifyResult{TypeName: "UnknownFailure"} +} + +//nolint:gocyclo // Rule-driven diagnosis intentionally keeps category handling in one switch. +func PerformDiagnosis(job *model.Job) DiagnosisResp { + resp := DiagnosisResp{JobName: job.JobName, Status: string(job.Status)} + result := CategorizeFailure(job) + resp.Category = result.TypeName + + switch result.TypeName { + case "OOMKilled": + resp.Diagnosis = "作业因内存溢出（OOM）被终止" + resp.Solution = "建议增加内存请求和限制，优化代码内存占用，并检查是否存在内存泄漏。" + resp.Confidence = diagnosisConfidenceHigh + resp.Severity = diagnosisSeverityCrit + case "ImagePullError": + resp.Diagnosis = "镜像拉取失败" + resp.Solution = "建议检查镜像名称、标签、仓库认证、网络连接，并确认镜像是否存在。" + resp.Confidence = diagnosisConfidenceHigh + resp.Severity = diagnosisSeverityError + case "SchedulingInsufficientResources": + resp.Diagnosis = "集群资源不足，无法调度" + resp.Solution = "建议降低资源请求，等待其他作业释放资源，或联系管理员扩容。" + resp.Confidence = diagnosisConfidenceHigh + resp.Severity = diagnosisSeverityWarn + case "SchedulingNodeSelectorMismatch": + resp.Diagnosis = "节点选择器不匹配" + resp.Solution = "建议检查节点标签配置，修改作业节点选择器，或联系管理员确认可用节点。" + resp.Confidence = diagnosisConfidenceHigh + resp.Severity = diagnosisSeverityError + case "SchedulingTaintMismatch": + resp.Diagnosis = "节点污点与容忍度不匹配" + resp.Solution = "建议检查节点 taint，在作业配置中补充 tolerations，或确认调度策略。" + resp.Confidence = diagnosisConfidenceHigh + resp.Severity = diagnosisSeverityError + case "SchedulingFailed": + resp.Diagnosis = "作业调度失败" + resp.Solution = "建议查看 FailedScheduling 事件原文，并检查资源请求、节点选择器与污点容忍。" + resp.Confidence = diagnosisConfidenceHigh + resp.Severity = diagnosisSeverityError + case "CrashLoopBackOff": + resp.Diagnosis = "容器持续崩溃重启" + resp.Solution = "建议查看容器日志，检查启动命令、配置文件和资源限制。" + resp.Confidence = diagnosisConfidenceHigh + resp.Severity = diagnosisSeverityCrit + case "VolumeMountFailed": + resp.Diagnosis = "存储卷挂载失败" + resp.Solution = "建议检查 PVC/PV 绑定状态、存储类、访问模式、挂载路径和权限。" + resp.Confidence = diagnosisConfidenceHigh + resp.Severity = diagnosisSeverityError + case "JobDeadlineExceeded": + resp.Diagnosis = "作业超出截止时间被终止" + resp.Solution = "建议评估并调大 activeDeadlineSeconds，或优化作业耗时。" + resp.Confidence = diagnosisConfidenceHigh + resp.Severity = diagnosisSeverityWarn + case "CommandNotFound": + resp.Diagnosis = "启动命令或文件不存在" + resp.Solution = "建议检查启动命令、镜像内目标文件/脚本、工作目录和 PATH 配置。" + resp.Confidence = diagnosisConfidenceHigh + resp.Severity = diagnosisSeverityError + case "GracefulTermination": + resp.Diagnosis = "作业收到终止信号并退出" + resp.Solution = "建议结合事件时间线确认是否为人工停止、调度回收或平台策略触发。" + resp.Confidence = diagnosisConfidenceHigh + resp.Severity = diagnosisSeverityWarn + case "Evicted": + resp.Diagnosis = "作业所在 Pod 被节点驱逐" + resp.Solution = "建议查看节点资源压力与驱逐原因，检查请求/限制是否合理。" + resp.Confidence = diagnosisConfidenceHigh + resp.Severity = diagnosisSeverityWarn + case "SegmentationFault": + resp.Diagnosis = "段错误（Segmentation Fault）" + resp.Solution = "建议检查非法内存访问、依赖兼容性、底层算子或 C/C++ 扩展代码。" + resp.Confidence = diagnosisConfidenceHigh + resp.Severity = diagnosisSeverityCrit + case "JobAbortedOrTerminated": + resp.Diagnosis = "作业被中止或终止" + resp.Solution = "建议检查是否有人工停止或控制器回收，并结合事件确认触发原因。" + resp.Confidence = diagnosisConfidenceHigh + resp.Severity = diagnosisSeverityWarn + default: + resp.Diagnosis = "未能自动诊断出具体原因" + resp.Solution = "建议查看作业日志和 Kubernetes 事件进行人工分析。" + resp.Confidence = "low" + resp.Severity = diagnosisSeverityError + } + + if job.TerminatedStates != nil && len(job.TerminatedStates.Data()) > 0 { + ts := job.TerminatedStates.Data()[0] + resp.Evidence.ExitCode = ts.ExitCode + resp.Evidence.ExitReason = ts.Reason + if title, suggestion, ok := exitCodeDiagnosis(ts.ExitCode); ok && + (resp.Category == "ContainerError" || resp.Category == "UnknownFailure") { + resp.Diagnosis = fmt.Sprintf("退出异常（Exit %d：%s）", ts.ExitCode, title) + resp.Solution = suggestion + resp.Confidence = "medium" + } + } + if job.Events != nil { + events := job.Events.Data() + for i := range events { + if events[i].Type == "Warning" || events[i].Type == "Error" { + resp.Evidence.Events = append(resp.Evidence.Events, events[i].Message) + } + } + if len(resp.Evidence.Events) > maxEvidenceEventsLimit { + resp.Evidence.Events = resp.Evidence.Events[:maxEvidenceEventsLimit] + } + } + return resp +} + +func sampleTerminated(ts *v1.ContainerStateTerminated) string { + if ts == nil { + return "" + } + return strings.TrimSpace(ts.Reason + " " + ts.Message) +} + +func sampleEvent(ev *v1.Event) string { + if ev == nil { + return "" + } + return strings.TrimSpace(ev.Message) +} + +func exitCodeDiagnosis(exitCode int32) (title, suggestion string, ok bool) { + mapping := map[int32]struct { + title string + suggestion string + }{ + 1: {"应用错误", "容器因应用程序错误而停止。建议优先查看日志中的错误堆栈，排查代码异常、依赖缺失、路径错误。"}, + 126: {"命令调用错误", "无法调用镜像中指定命令。请确认命令路径正确，且具备可执行权限。"}, + 127: {"命令或文件不存在", "找不到镜像中指定命令或文件。请检查启动命令、工作目录、文件路径与镜像内容是否一致。"}, + 137: {"立即终止（SIGKILL）", "通常为内存不足或被系统强制终止。建议增加内存申请并结合日志确认是否发生 OOM。"}, + 139: {"分段错误（SIGSEGV）", "进程发生非法内存访问。请检查依赖兼容性、底层算子或 C/C++ 扩展代码。"}, + 143: {"优雅终止（SIGTERM）", "进程收到终止信号后退出，可能是调度或人工停止触发。请结合事件时间线判断是否为预期行为。"}, + } + item, exists := mapping[exitCode] + if !exists { + return "", "", false + } + return item.title, item.suggestion, true +} diff --git a/backend/internal/handler/context.go b/backend/internal/handler/context.go index af5d27b9e..76aac1659 100644 --- a/backend/internal/handler/context.go +++ b/backend/internal/handler/context.go @@ -17,6 +17,7 @@ import ( "github.com/raids-lab/crater/dao/model" "github.com/raids-lab/crater/dao/query" + "github.com/raids-lab/crater/internal/bizerr" "github.com/raids-lab/crater/internal/resputil" "github.com/raids-lab/crater/internal/service" "github.com/raids-lab/crater/internal/util" @@ -53,12 +54,16 @@ func (mgr *ContextMgr) GetName() string { return mgr.name } func (mgr *ContextMgr) RegisterPublic(_ *gin.RouterGroup) {} +//nolint:dupl // Router registration blocks naturally look similar across managers. func (mgr *ContextMgr) RegisterProtected(g *gin.RouterGroup) { g.GET("prequeue", mgr.GetPrequeueStatus) g.GET("quota", mgr.GetQuota) g.GET("job-resource-summary", mgr.GetJobResourceSummary) g.POST("resource-limit-check", mgr.CheckResourceLimit) g.GET("billing/summary", mgr.GetBillingSummary) + g.GET("llm", mgr.GetUserLLMConfig) + g.PUT("llm", mgr.UpdateUserLLMConfig) + g.DELETE("llm", mgr.ResetUserLLMConfig) g.PUT("attributes", mgr.UpdateUserAttributes) g.POST("email/code", mgr.SendUserVerificationCode) g.POST("email/update", mgr.UpdateUserEmail) @@ -100,6 +105,15 @@ type ( Memory JobResourceSummaryUsageResp `json:"memory"` Accelerators []JobResourceSummaryAcceleratorResp `json:"accelerators"` } + UserLLMConfigResp struct { + BaseURL string `json:"baseUrl"` + APIKey string `json:"apiKey"` + ModelName string `json:"modelName"` + Source string `json:"source"` + UsingPersonal bool `json:"usingPersonal"` + Complete bool `json:"complete"` + HasAPIKey bool `json:"hasApiKey"` + } ) // GetPrequeueStatus godoc @@ -292,6 +306,75 @@ func hasPositiveQuantity(value string) bool { return quantity.MilliValue() > 0 } +func (mgr *ContextMgr) GetUserLLMConfig(c *gin.Context) { + if mgr.configService == nil { + resputil.HandleError(c, bizerr.Internal.ServiceError.New("config service is not initialized")) + return + } + token := util.GetToken(c) + status, err := mgr.configService.GetEffectiveLLMConfigStatus(c.Request.Context(), token.UserID) + if err != nil { + resputil.HandleError(c, bizerr.Internal.ServiceError.Wrap(err, "failed to load LLM config")) + return + } + + resp := UserLLMConfigResp{ + Source: status.Source, + UsingPersonal: status.UsingPersonal, + Complete: status.Complete, + HasAPIKey: status.HasAPIKey, + } + if status.Config != nil { + resp.BaseURL = strings.TrimSpace(status.Config.BaseURL) + resp.ModelName = strings.TrimSpace(status.Config.ModelName) + if status.UsingPersonal && strings.TrimSpace(status.Config.APIKey) != "" { + resp.APIKey = service.MaskedAPIKeyPlaceholder + } + } + resputil.Success(c, resp) +} + +func (mgr *ContextMgr) UpdateUserLLMConfig(c *gin.Context) { + if mgr.configService == nil { + resputil.HandleError(c, bizerr.Internal.ServiceError.New("config service is not initialized")) + return + } + var req UpdateLLMConfigReq + if err := c.ShouldBindJSON(&req); err != nil { + resputil.HandleError(c, bizerr.BadRequest.ParameterError.New(err.Error())) + return + } + + token := util.GetToken(c) + serviceCfg := &service.LLMConfig{ + BaseURL: req.BaseURL, + APIKey: req.APIKey, + ModelName: req.ModelName, + } + if err := mgr.configService.UpdateUserLLMConfig(c.Request.Context(), token.UserID, serviceCfg, req.Validate); err != nil { + if strings.Contains(err.Error(), "validation failed") { + resputil.HandleError(c, bizerr.Conflict.ResourceStatusError.New("LLM connection check failed. Please verify your settings.")) + return + } + resputil.HandleError(c, err) + return + } + resputil.Success(c, "User LLM configuration updated successfully") +} + +func (mgr *ContextMgr) ResetUserLLMConfig(c *gin.Context) { + if mgr.configService == nil { + resputil.HandleError(c, bizerr.Internal.ServiceError.New("config service is not initialized")) + return + } + token := util.GetToken(c) + if err := mgr.configService.ResetUserLLMConfig(c.Request.Context(), token.UserID); err != nil { + resputil.HandleError(c, err) + return + } + resputil.Success(c, "User LLM configuration reset successfully") +} + type ( UserInfoResp struct { ID uint `json:"id"` diff --git a/backend/internal/handler/dataset.go b/backend/internal/handler/dataset.go index dc8f1cd42..59c254776 100644 --- a/backend/internal/handler/dataset.go +++ b/backend/internal/handler/dataset.go @@ -1,9 +1,7 @@ package handler import ( - "crypto/sha256" "fmt" - "net/http" "path/filepath" "regexp" "strings" @@ -16,14 +14,11 @@ import ( "github.com/raids-lab/crater/dao/model" "github.com/raids-lab/crater/dao/query" "github.com/raids-lab/crater/internal/bizerr" - "github.com/raids-lab/crater/internal/governance/modeldataset" "github.com/raids-lab/crater/internal/resputil" "github.com/raids-lab/crater/internal/util" "github.com/raids-lab/crater/pkg/config" ) -const autoDownloadTag = "auto-download" - //nolint:gochecknoinits // This is the standard way to register a gin handler. func init() { Registers = append(Registers, NewDatasetMgr) @@ -41,9 +36,7 @@ func NewDatasetMgr(_ *RegisterConfig) Manager { func (mgr *DatasetMgr) GetName() string { return mgr.name } -func (mgr *DatasetMgr) RegisterPublic(g *gin.RouterGroup) { - g.GET("/source-logo/:sourceId", mgr.GetSourceLogo) -} +func (mgr *DatasetMgr) RegisterPublic(_ *gin.RouterGroup) {} func (mgr *DatasetMgr) RegisterProtected(g *gin.RouterGroup) { g.GET("/mydataset", mgr.GetDatasets) @@ -69,47 +62,6 @@ func (mgr *DatasetMgr) RegisterAdmin(g *gin.RouterGroup) { g.POST("/cancelshare/queue", mgr.AdmincancelShareDatasetWithQueue) } -// GetSourceLogo returns a platform-cached public source logo. -// -// @Summary 获取平台缓存的模型或数据集来源 Logo -// @Description 返回平台缓存的来源 Logo，不要求浏览器访问外部模型站点 -// @Tags Dataset -// @Produce image/png,image/jpeg,image/webp,image/svg+xml -// @Param sourceId path int true "来源 ID" -// @Success 200 {file} binary -// @Failure 400 -// @Failure 404 -// @Router /dataset/source-logo/{sourceId} [get] -func (mgr *DatasetMgr) GetSourceLogo(c *gin.Context) { - var req struct { - SourceID uint `uri:"sourceId" binding:"required"` - } - if err := c.ShouldBindUri(&req); err != nil { - c.Status(http.StatusBadRequest) - return - } - var source model.ModelDatasetSource - if err := query.GetDB().WithContext(c). - Select("logo_data", "logo_content_type", "updated_at"). - First(&source, req.SourceID).Error; err != nil || len(source.LogoData) == 0 { - c.Status(http.StatusNotFound) - return - } - contentType := source.LogoContentType - if contentType == "" { - contentType = http.DetectContentType(source.LogoData) - } - etag := fmt.Sprintf("\"%x\"", sha256.Sum256(source.LogoData)) - c.Header("Cache-Control", "public, no-cache") - c.Header("ETag", etag) - c.Header("X-Content-Type-Options", "nosniff") - if c.GetHeader("If-None-Match") == etag { - c.Status(http.StatusNotModified) - return - } - c.Data(http.StatusOK, contentType, source.LogoData) -} - type DatasetResp struct { Name string `json:"name"` ID uint `json:"id"` @@ -865,7 +817,7 @@ func (mgr *DatasetMgr) UpdateDataset(c *gin.Context) { tempExtra := dataset.Extra.Data() tempExtra.Tags = make([]string, 0, len(req.Tags)) for _, tag := range req.Tags { - if tag != autoDownloadTag { + if tag != "auto-download" { tempExtra.Tags = append(tempExtra.Tags, tag) } } @@ -1111,35 +1063,6 @@ func listDatasetResponses(c *gin.Context, ids []uint) ([]DatasetResp, error) { return convertDatasetBatch(c, datasets) } -// deduplicateDownloadedDatasets keeps one canonical row for a downloaded -// logical resource. It deliberately does not merge user-created resources, -// which may legitimately reuse a name at different storage paths. Prefer the -// lowest ID because it is the original row and is most likely to carry existing -// sharing associations. -func deduplicateDownloadedDatasets( - datasets []*model.Dataset, downloadedResourceKeys map[string]struct{}, -) []*model.Dataset { - result := make([]*model.Dataset, 0, len(datasets)) - positions := make(map[string]int, len(datasets)) - for _, dataset := range datasets { - key := string(dataset.Type) + "\x00" + strings.ToLower(strings.TrimSpace(dataset.Name)) - if _, downloaded := downloadedResourceKeys[key]; !downloaded { - result = append(result, dataset) - continue - } - position, exists := positions[key] - if !exists { - positions[key] = len(result) - result = append(result, dataset) - continue - } - if dataset.ID < result[position].ID { - result[position] = dataset - } - } - return result -} - func convertDataset(c *gin.Context, dataset *model.Dataset) DatasetResp { responses, err := convertDatasetBatch(c, []*model.Dataset{dataset}) if err != nil || len(responses) == 0 { @@ -1148,7 +1071,6 @@ func convertDataset(c *gin.Context, dataset *model.Dataset) DatasetResp { return responses[0] } -//nolint:gocyclo // Batch enrichment deliberately keeps legacy and source-table fallback in one pass. func convertDatasetBatch(c *gin.Context, datasets []*model.Dataset) ([]DatasetResp, error) { names := make([]string, 0, len(datasets)) seenNames := make(map[string]struct{}, len(datasets)) @@ -1164,7 +1086,6 @@ func convertDatasetBatch(c *gin.Context, datasets []*model.Dataset) ([]DatasetRe } downloadsByResource := make(map[string]*model.ModelDownload, len(names)) - downloadedResourceKeys := make(map[string]struct{}, len(names)) if len(names) > 0 { q := query.ModelDownload downloads, err := q.WithContext(c). @@ -1179,174 +1100,20 @@ func convertDatasetBatch(c *gin.Context, datasets []*model.Dataset) ([]DatasetRe if _, exists := downloadsByResource[key]; !exists { downloadsByResource[key] = download } - downloadedResourceKeys[resourceMetadataKey( - strings.ToLower(strings.TrimSpace(download.Name)), string(download.Category), - )] = struct{}{} - } - } - datasets = deduplicateDownloadedDatasets(datasets, downloadedResourceKeys) - sourceIDs := make([]uint, 0, len(datasets)) - for _, dataset := range datasets { - if dataset.ModelDatasetSourceID != nil { - sourceIDs = append(sourceIDs, *dataset.ModelDatasetSourceID) - } - } - sourcesByID := make(map[uint]*model.ModelDatasetSource, len(sourceIDs)) - if len(sourceIDs) > 0 { - var sources []*model.ModelDatasetSource - if err := query.GetDB().WithContext(c).Where("id IN ?", sourceIDs).Find(&sources).Error; err != nil { - return nil, err - } - for _, source := range sources { - sourcesByID[source.ID] = source } } - logoKeys := make(map[string]struct{}, len(datasets)+len(downloadsByResource)) - for _, dataset := range datasets { - if key := normalizedOrganizationLogoKey("", dataset.Name); key != "" { - logoKeys[key] = struct{}{} - } - } - for _, download := range downloadsByResource { - if key := downloadOrganizationLogoKey(download); key != "" { - logoKeys[key] = struct{}{} - } - } - logoSourceIDs, err := loadOrganizationLogoSourceIDs(c, sourcesByID, logoKeys) - if err != nil { - return nil, err - } responses := make([]DatasetResp, 0, len(datasets)) for _, dataset := range datasets { resp := baseDatasetResp(dataset) - if dataset.ModelDatasetSourceID != nil && sourcesByID[*dataset.ModelDatasetSourceID] != nil { - source := sourcesByID[*dataset.ModelDatasetSourceID] - enrichDatasetRespFromSource(&resp, source, logoSourceIDs[organizationLogoKey(source)]) - } else if download := downloadsByResource[resourceMetadataKey(dataset.Name, string(dataset.Type))]; download != nil { - enrichDatasetResp(&resp, download, logoSourceIDs[downloadOrganizationLogoKey(download)]) - } else if logoSourceID := logoSourceIDs[normalizedOrganizationLogoKey("", dataset.Name)]; logoSourceID != 0 { - resp.Organization = strings.SplitN(dataset.Name, "/", 2)[0] - resp.OrganizationURL = fmt.Sprintf("/api/dataset/source-logo/%d", logoSourceID) + if download := downloadsByResource[resourceMetadataKey(dataset.Name, string(dataset.Type))]; download != nil { + enrichDatasetResp(&resp, download) } responses = append(responses, resp) } return responses, nil } -func organizationLogoKey(source *model.ModelDatasetSource) string { - return normalizedOrganizationLogoKey(source.Organization, source.RepositoryID) -} - -func downloadOrganizationLogoKey(download *model.ModelDownload) string { - return normalizedOrganizationLogoKey(download.Organization, download.Name) -} - -func normalizedOrganizationLogoKey(organization, repositoryID string) string { - organization = strings.TrimSpace(organization) - if organization == "" { - parts := strings.SplitN(strings.TrimSpace(repositoryID), "/", 2) - if len(parts) < 2 { - return "" - } - organization = parts[0] - } - return strings.ToLower(organization) -} - -// loadOrganizationLogoSourceIDs lets repositories from the same organization reuse one -// platform-cached logo. Historical source rows may contain LogoURL but no LogoData because -// they were refreshed before binary logo caching was enabled. -func loadOrganizationLogoSourceIDs( - c *gin.Context, - sourcesByID map[uint]*model.ModelDatasetSource, - additionalKeys map[string]struct{}, -) (map[string]uint, error) { - result := make(map[string]uint) - missingKeys := make(map[string]struct{}) - for _, source := range sourcesByID { - key := organizationLogoKey(source) - if key == "" { - continue - } - if len(source.LogoData) > 0 { - result[key] = source.ID - } else { - missingKeys[key] = struct{}{} - } - } - for key := range additionalKeys { - if key != "" && result[key] == 0 { - missingKeys[key] = struct{}{} - } - } - if len(missingKeys) == 0 { - return result, nil - } - - organizations := make([]string, 0, len(missingKeys)) - for key := range missingKeys { - if result[key] == 0 { - organizations = append(organizations, key) - } - } - if len(organizations) == 0 { - return result, nil - } - - var cachedSources []model.ModelDatasetSource - if err := query.GetDB().WithContext(c). - Where("LOWER(organization) IN ? AND octet_length(logo_data) > 0", organizations). - Order("updated_at DESC"). - Find(&cachedSources).Error; err != nil { - return nil, err - } - for i := range cachedSources { - key := organizationLogoKey(&cachedSources[i]) - if result[key] == 0 { - result[key] = cachedSources[i].ID - } - } - return result, nil -} - -func enrichDatasetRespFromSource(resp *DatasetResp, source *model.ModelDatasetSource, logoSourceID uint) { - resp.DownloadCount = int(source.Downloads) - resp.Likes = source.Likes - resp.Source = string(source.Provider) - resp.Organization = source.Organization - if resp.Organization == "" { - resp.Organization = strings.SplitN(source.RepositoryID, "/", 2)[0] - } - if logoSourceID != 0 { - resp.OrganizationURL = fmt.Sprintf("/api/dataset/source-logo/%d", logoSourceID) - } - resp.DisplayName = source.DisplayName - resp.Readme = source.Readme - resp.License = source.License - resp.Task = source.Task - resp.Library = source.Library - resp.ModelType = source.ModelType - resp.ParameterCount = source.ParameterCount - resp.SourcePrivate = source.Private - resp.SourceGated = source.Gated - resp.LoginRequired = source.LoginRequired - resp.SourceCreatedAt = source.SourceCreatedAt - resp.SourceUpdatedAt = source.SourceUpdatedAt - extra := resp.Extra.Data() - filteredTags := extra.Tags[:0] - for _, tag := range extra.Tags { - if tag != autoDownloadTag { - filteredTags = append(filteredTags, tag) - } - } - extra.Tags = filteredTags - if extra.WebURL == nil && source.RepositoryURL != "" { - extra.WebURL = &source.RepositoryURL - } - resp.Extra = datatypes.NewJSONType(extra) -} - func resourceMetadataKey(name, category string) string { return category + "\x00" + name } @@ -1355,7 +1122,7 @@ func baseDatasetResp(dataset *model.Dataset) DatasetResp { return DatasetResp{ Name: dataset.Name, Describe: dataset.Describe, - URL: logicalDatasetStoragePath(dataset.URL), + URL: dataset.URL, ID: dataset.ID, CreatedAt: dataset.CreatedAt, UpdatedAt: dataset.UpdatedAt, @@ -1370,25 +1137,7 @@ func baseDatasetResp(dataset *model.Dataset) DatasetResp { } } -func logicalDatasetStoragePath(rawURL string) string { - prefixes := config.GetConfig().Storage.Prefix - mappings := []struct { - physical string - logical string - }{ - {physical: prefixes.Public, logical: model.PublicPath}, - {physical: prefixes.Account, logical: model.AccountPath}, - {physical: prefixes.User, logical: model.UserPath}, - } - for _, mapping := range mappings { - if logical, matched := modeldataset.LogicalStoragePath(rawURL, mapping.logical, mapping.physical); matched { - return logical - } - } - return rawURL -} - -func enrichDatasetResp(resp *DatasetResp, download *model.ModelDownload, logoSourceID uint) { +func enrichDatasetResp(resp *DatasetResp, download *model.ModelDownload) { if resp.SizeBytes == 0 { resp.SizeBytes = download.SizeBytes } @@ -1401,7 +1150,7 @@ func enrichDatasetResp(resp *DatasetResp, download *model.ModelDownload, logoSou extra := resp.Extra.Data() filteredTags := extra.Tags[:0] for _, tag := range extra.Tags { - if tag != autoDownloadTag { + if tag != "auto-download" { filteredTags = append(filteredTags, tag) } } @@ -1412,9 +1161,7 @@ func enrichDatasetResp(resp *DatasetResp, download *model.ModelDownload, logoSou } resp.Extra = datatypes.NewJSONType(extra) resp.Organization = download.Organization - if logoSourceID != 0 { - resp.OrganizationURL = fmt.Sprintf("/api/dataset/source-logo/%d", logoSourceID) - } + resp.OrganizationURL = download.LogoURL resp.DisplayName = download.DisplayName resp.Readme = download.SourceReadme resp.License = download.License diff --git a/backend/internal/handler/dataset_test.go b/backend/internal/handler/dataset_test.go deleted file mode 100644 index 366299957..000000000 --- a/backend/internal/handler/dataset_test.go +++ /dev/null @@ -1,72 +0,0 @@ -package handler - -import ( - "testing" - - "github.com/raids-lab/crater/dao/model" -) - -func TestNormalizedOrganizationLogoKey(t *testing.T) { - t.Parallel() - - tests := []struct { - name string - organization string - repositoryID string - want string - }{ - {name: "explicit organization", organization: " Qwen ", repositoryID: "ignored", want: "qwen"}, - {name: "legacy repository", repositoryID: "qwen/qwen2-0.5b", want: "qwen"}, - {name: "manual resource name", repositoryID: "qwen2-0.5b", want: ""}, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - if got := normalizedOrganizationLogoKey(tt.organization, tt.repositoryID); got != tt.want { - t.Fatalf("normalizedOrganizationLogoKey(%q, %q) = %q, want %q", tt.organization, tt.repositoryID, got, tt.want) - } - }) - } -} - -func TestDeduplicateDatasets(t *testing.T) { - t.Parallel() - - newerDuplicate := &model.Dataset{Name: "SKYLENAGE/SkyJM-Gen-4B", Type: model.DataTypeModel} - newerDuplicate.ID = 92 - original := &model.Dataset{Name: "skylenage/SkyJM-Gen-4B", Type: model.DataTypeModel} - original.ID = 91 - datasetWithSameName := &model.Dataset{Name: "SKYLENAGE/SkyJM-Gen-4B", Type: model.DataTypeDataset} - datasetWithSameName.ID = 93 - - downloadedKeys := map[string]struct{}{ - resourceMetadataKey("skylenage/skyjm-gen-4b", string(model.DataTypeModel)): {}, - } - got := deduplicateDownloadedDatasets( - []*model.Dataset{newerDuplicate, original, datasetWithSameName}, downloadedKeys, - ) - if len(got) != 2 { - t.Fatalf("deduplicateDatasets() returned %d rows, want 2", len(got)) - } - if got[0].ID != original.ID { - t.Fatalf("deduplicateDatasets() kept ID %d, want canonical ID %d", got[0].ID, original.ID) - } - if got[1].ID != datasetWithSameName.ID { - t.Fatalf("deduplicateDatasets() incorrectly merged different resource types") - } -} - -func TestDeduplicateDownloadedDatasetsPreservesUserResources(t *testing.T) { - t.Parallel() - - first := &model.Dataset{Name: "experiment", Type: model.DataTypeDataset} - first.ID = 1 - second := &model.Dataset{Name: "experiment", Type: model.DataTypeDataset} - second.ID = 2 - - got := deduplicateDownloadedDatasets([]*model.Dataset{first, second}, nil) - if len(got) != 2 { - t.Fatalf("deduplicateDownloadedDatasets() merged user resources, got %d rows", len(got)) - } -} diff --git a/backend/internal/handler/image/agent_insights.go b/backend/internal/handler/image/agent_insights.go new file mode 100644 index 000000000..43283a9a0 --- /dev/null +++ b/backend/internal/handler/image/agent_insights.go @@ -0,0 +1,188 @@ +package image + +import ( + "context" + "sort" + + "github.com/raids-lab/crater/dao/model" + "github.com/raids-lab/crater/dao/query" + "github.com/raids-lab/crater/internal/handler" + "github.com/raids-lab/crater/internal/util" +) + +//nolint:gochecknoinits // The agent package discovers this facade through handler registration. +func init() { + handler.RegisterImageInsightReaderFactory(NewAgentImageInsightReader) +} + +type agentImageInsightReader struct{} + +func NewAgentImageInsightReader(_ *handler.RegisterConfig) handler.ImageInsightReader { + return &agentImageInsightReader{} +} + +func (r *agentImageInsightReader) ListAccessibleImages( + ctx context.Context, + token util.JWTMessage, +) ([]handler.ImageAccessRecord, error) { + if token.RolePlatform == model.RoleAdmin { + return r.listAdminImages(ctx) + } + return r.listUserImages(ctx, token) +} + +func (r *agentImageInsightReader) listAdminImages(ctx context.Context) ([]handler.ImageAccessRecord, error) { + publicImageIDs, err := r.listPublicImageIDs(ctx) + if err != nil { + return nil, err + } + publicIDSet := make(map[uint]struct{}, len(publicImageIDs)) + for _, id := range publicImageIDs { + publicIDSet[id] = struct{}{} + } + + iq := query.Image + images, err := iq.WithContext(ctx).Preload(iq.User).Order(iq.CreatedAt.Desc()).Find() + if err != nil { + return nil, err + } + + records := make([]handler.ImageAccessRecord, 0, len(images)) + for _, item := range images { + status := model.Private + if _, ok := publicIDSet[item.ID]; ok || item.IsPublic { + status = model.Public + } + records = append(records, handler.ImageAccessRecord{ + Image: item, + ShareStatus: status, + }) + } + return records, nil +} + +func (r *agentImageInsightReader) listUserImages( + ctx context.Context, + token util.JWTMessage, +) ([]handler.ImageAccessRecord, error) { + results := make([]handler.ImageAccessRecord, 0) + seen := make(map[uint]struct{}) + appendUnique := func(images []*model.Image, shareStatus model.ImageShareType) { + for _, imageRecord := range images { + if imageRecord == nil || imageRecord.ID == 0 || imageRecord.ImageLink == "" { + continue + } + if _, ok := seen[imageRecord.ID]; ok { + continue + } + seen[imageRecord.ID] = struct{}{} + results = append(results, handler.ImageAccessRecord{ + Image: imageRecord, + ShareStatus: shareStatus, + }) + } + } + + iq := query.Image + oldPublicImages, err := iq.WithContext(ctx). + Preload(iq.User). + Where(iq.IsPublic.Is(true)). + Order(iq.CreatedAt.Desc()). + Find() + if err != nil { + return nil, err + } + appendUnique(oldPublicImages, model.Public) + + newPublicImages, err := r.listAccountSharedImages(ctx, model.DefaultAccountID) + if err != nil { + return nil, err + } + appendUnique(newPublicImages, model.Public) + + accountImages, err := r.listAccountSharedImages(ctx, token.AccountID) + if err != nil { + return nil, err + } + appendUnique(accountImages, model.AccountShare) + + privateImages, err := iq.WithContext(ctx). + Preload(iq.User). + Where(iq.UserID.Eq(token.UserID)). + Order(iq.CreatedAt.Desc()). + Find() + if err != nil { + return nil, err + } + appendUnique(privateImages, model.Private) + + userImages, err := r.listUserSharedImages(ctx, token.UserID) + if err != nil { + return nil, err + } + appendUnique(userImages, model.UserShare) + + sort.Slice(results, func(i, j int) bool { + return results[i].Image.CreatedAt.After(results[j].Image.CreatedAt) + }) + return results, nil +} + +func (r *agentImageInsightReader) listPublicImageIDs(ctx context.Context) ([]uint, error) { + imageAccountQuery := query.ImageAccount + publicImageIDs := []uint{} + if err := imageAccountQuery.WithContext(ctx). + Where(imageAccountQuery.AccountID.Eq(model.DefaultAccountID)). + Pluck(imageAccountQuery.ImageID, &publicImageIDs); err != nil { + return nil, err + } + + imageQuery := query.Image + oldPublicImageIDs := []uint{} + if err := imageQuery.WithContext(ctx). + Where(imageQuery.IsPublic.Is(true)). + Pluck(imageQuery.ID, &oldPublicImageIDs); err != nil { + return nil, err + } + return append(publicImageIDs, oldPublicImageIDs...), nil +} + +//nolint:dupl // GORM gen exposes separate typed query builders for account and user image shares. +func (r *agentImageInsightReader) listAccountSharedImages(ctx context.Context, accountID uint) ([]*model.Image, error) { + imageShareQuery := query.ImageAccount + imageShares, err := imageShareQuery.WithContext(ctx). + Preload(imageShareQuery.Image). + Preload(imageShareQuery.Image.User). + Where(imageShareQuery.AccountID.Eq(accountID)). + Find() + if err != nil { + return nil, err + } + return imagesFromShares(imageShares, func(imageShare *model.ImageAccount) *model.Image { + return &imageShare.Image + }), nil +} + +//nolint:dupl // GORM gen exposes separate typed query builders for account and user image shares. +func (r *agentImageInsightReader) listUserSharedImages(ctx context.Context, userID uint) ([]*model.Image, error) { + imageShareQuery := query.ImageUser + imageShares, err := imageShareQuery.WithContext(ctx). + Preload(imageShareQuery.Image). + Preload(imageShareQuery.Image.User). + Where(imageShareQuery.UserID.Eq(userID)). + Find() + if err != nil { + return nil, err + } + return imagesFromShares(imageShares, func(imageShare *model.ImageUser) *model.Image { + return &imageShare.Image + }), nil +} + +func imagesFromShares[T any](imageShares []T, imageOf func(T) *model.Image) []*model.Image { + images := make([]*model.Image, 0, len(imageShares)) + for _, imageShare := range imageShares { + images = append(images, imageOf(imageShare)) + } + return images +} diff --git a/backend/internal/handler/image/buildmgr.go b/backend/internal/handler/image/buildmgr.go index fa2c9cf85..8363fbbed 100644 --- a/backend/internal/handler/image/buildmgr.go +++ b/backend/internal/handler/image/buildmgr.go @@ -1,17 +1,14 @@ package image import ( - "errors" "fmt" "github.com/gin-gonic/gin" - "gorm.io/gorm" corev1 "k8s.io/api/core/v1" "k8s.io/klog/v2" "github.com/raids-lab/crater/dao/model" "github.com/raids-lab/crater/dao/query" - "github.com/raids-lab/crater/internal/bizerr" "github.com/raids-lab/crater/internal/resputil" "github.com/raids-lab/crater/internal/util" ) @@ -213,14 +210,18 @@ func (mgr *ImagePackMgr) deleteKanikoByID(c *gin.Context, isAdminMode bool, kani // @Param name query string true "获取ImagePack的name" // @Router /v1/images/getbyname [GET] func (mgr *ImagePackMgr) GetKanikoByImagePackName(c *gin.Context) { + kanikoQuery := query.Kaniko var req GetKanikoRequest - if err := c.ShouldBindQuery(&req); err != nil { + var err error + if err = c.ShouldBindQuery(&req); err != nil { msg := fmt.Sprintf("validate image get parameters failed, err %v", err) resputil.BadRequestError(c, msg) return } - kaniko, err := mgr.findCurrentUserKaniko(c, req.ImagePackName, 0) - if err != nil { + var kaniko *model.Kaniko + if kaniko, err = kanikoQuery.WithContext(c). + Where(kanikoQuery.ImagePackName.Eq(req.ImagePackName)). + First(); err != nil { msg := fmt.Sprintf("fetch kaniko by name failed, err %v", err) resputil.BadRequestError(c, msg) return @@ -254,14 +255,18 @@ func (mgr *ImagePackMgr) GetKanikoByImagePackName(c *gin.Context) { // @Param name query string true "获取ImagePack的name" // @Router /v1/images/get [GET] func (mgr *ImagePackMgr) GetKanikoTemplateByImagePackName(c *gin.Context) { + kanikoQuery := query.Kaniko var req GetKanikoRequest - if err := c.ShouldBindQuery(&req); err != nil { + var err error + if err = c.ShouldBindQuery(&req); err != nil { msg := fmt.Sprintf("validate image get parameters failed, err %v", err) resputil.BadRequestError(c, msg) return } - kaniko, err := mgr.findCurrentUserKaniko(c, req.ImagePackName, 0) - if err != nil { + var kaniko *model.Kaniko + if kaniko, err = kanikoQuery.WithContext(c). + Where(kanikoQuery.ImagePackName.Eq(req.ImagePackName)). + First(); err != nil { msg := fmt.Sprintf("fetch kaniko by name failed, err %v", err) resputil.BadRequestError(c, msg) return @@ -288,15 +293,6 @@ func (mgr *ImagePackMgr) GetImagepackPodName(c *gin.Context) { resputil.BadRequestError(c, msg) return } - if _, err = mgr.findCurrentUserKaniko(c, "", req.ID); err != nil { - if errors.Is(err, gorm.ErrRecordNotFound) { - resputil.HandleError(c, bizerr.Forbidden.PermissionDenied.New("permission denied to view this image build")) - return - } - klog.Errorf("fetch kaniko pod failed, id: %d, err: %v", req.ID, err) - resputil.HandleError(c, bizerr.Internal.DatabaseError.Wrap(err, "fetch kaniko pod failed")) - return - } podName, podNameSpace, nodeName := mgr.getPodName(c, req.ID) resp := GetKanikoPodResponse{ PodName: podName, @@ -306,17 +302,6 @@ func (mgr *ImagePackMgr) GetImagepackPodName(c *gin.Context) { resputil.Success(c, resp) } -func (mgr *ImagePackMgr) findCurrentUserKaniko(c *gin.Context, imagePackName string, id uint) (*model.Kaniko, error) { - kanikoQuery := query.Kaniko - q := kanikoQuery.WithContext(c).Where(kanikoQuery.UserID.Eq(util.GetToken(c).UserID)) - if imagePackName != "" { - q = q.Where(kanikoQuery.ImagePackName.Eq(imagePackName)) - } else { - q = q.Where(kanikoQuery.ID.Eq(id)) - } - return q.First() -} - func (mgr *ImagePackMgr) getPodName(c *gin.Context, kanikoID uint) (name, ns, nodeName string) { var kaniko *model.Kaniko var err error diff --git a/backend/internal/handler/image/harbor.go b/backend/internal/handler/image/harbor.go index c5781b468..8398ced5e 100644 --- a/backend/internal/handler/image/harbor.go +++ b/backend/internal/handler/image/harbor.go @@ -132,6 +132,7 @@ func (mgr *ImagePackMgr) UserGetProjectCredential(c *gin.Context) { resputil.Error(c, "add project member failed", resputil.NotSpecified) return } + fmt.Printf("username: %s, password: %s\n", token.Username, password) resp := GetProjectCredentialResponse{ Name: &token.Username, Password: &password, diff --git a/backend/internal/handler/image/image.go b/backend/internal/handler/image/image.go index 27f8f8c87..11ba6e149 100644 --- a/backend/internal/handler/image/image.go +++ b/backend/internal/handler/image/image.go @@ -1,7 +1,6 @@ package image import ( - "errors" "fmt" "sort" @@ -9,11 +8,9 @@ import ( "k8s.io/klog/v2" "gorm.io/datatypes" - "gorm.io/gorm" "github.com/raids-lab/crater/dao/model" "github.com/raids-lab/crater/dao/query" - "github.com/raids-lab/crater/internal/bizerr" "github.com/raids-lab/crater/internal/resputil" "github.com/raids-lab/crater/internal/util" ) @@ -283,13 +280,7 @@ func (mgr *ImagePackMgr) DeleteImageByID(c *gin.Context) { } imageID := deleteImageRequest.ID imageQuery := query.Image - if !mgr.requireImageOwner(c, imageID) { - return - } - if _, err = imageQuery.WithContext(c). - Where(imageQuery.ID.Eq(imageID)). - Where(imageQuery.UserID.Eq(util.GetToken(c).UserID)). - Delete(); err != nil { + if _, err = imageQuery.WithContext(c).Where(imageQuery.ID.Eq(imageID)).Delete(); err != nil { klog.Errorf("delete image entity failed! err:%v", err) resputil.Error(c, "failed to delete image", resputil.NotSpecified) } @@ -314,11 +305,6 @@ func (mgr *ImagePackMgr) UserDeleteImageByIDList(c *gin.Context) { resputil.BadRequestError(c, msg) return } - for _, imageID := range deleteImageListRequest.IDList { - if !mgr.requireImageOwner(c, imageID) { - return - } - } flag := mgr.deleteImageByIDList(c, false, deleteImageListRequest.IDList) if flag { resputil.Success(c, "") @@ -437,9 +423,6 @@ func (mgr *ImagePackMgr) UserChangeImageDescription(c *gin.Context) { resputil.BadRequestError(c, "validate failed") return } - if !mgr.requireImageOwner(c, req.ID) { - return - } mgr.changeImageDescription(c, false, req.ID, req.Description) } @@ -493,9 +476,6 @@ func (mgr *ImagePackMgr) UserChangeImageTaskType(c *gin.Context) { resputil.BadRequestError(c, "validate failed") return } - if !mgr.requireImageOwner(c, req.ID) { - return - } mgr.changeImageTaskType(c, false, req.ID, req.TaskType) } @@ -597,10 +577,7 @@ func (mgr *ImagePackMgr) UserChangeImageTags(c *gin.Context) { resputil.BadRequestError(c, "validate failed") return } - if !mgr.requireImageOwner(c, req.ID) { - return - } - mgr.changeImageTags(c, false, req.ID, req.Tags) + mgr.changeImageTags(c, req.ID, req.Tags) } // AdminChangeImageTagsType godoc @@ -619,16 +596,13 @@ func (mgr *ImagePackMgr) AdminChangeImageTags(c *gin.Context) { resputil.BadRequestError(c, "validate failed") return } - mgr.changeImageTags(c, true, req.ID, req.Tags) + mgr.changeImageTags(c, req.ID, req.Tags) } -func (mgr *ImagePackMgr) changeImageTags(c *gin.Context, isAdminMode bool, imageID uint, newTags []string) { +func (mgr *ImagePackMgr) changeImageTags(c *gin.Context, imageID uint, newTags []string) { imageQuery := query.Image - specifiedQuery := imageQuery.WithContext(c) - if !isAdminMode { - specifiedQuery = specifiedQuery.Where(imageQuery.UserID.Eq(util.GetToken(c).UserID)) - } - if _, err := specifiedQuery.Where(imageQuery.ID.Eq(imageID)). + if _, err := imageQuery.WithContext(c). + Where(imageQuery.ID.Eq(imageID)). Update(imageQuery.Tags, datatypes.NewJSONType(newTags)); err != nil { klog.Errorf("update image tags failed, err %v", err) resputil.BadRequestError(c, "update tags failed") @@ -652,9 +626,6 @@ func (mgr *ImagePackMgr) UserShareImage(c *gin.Context) { resputil.BadRequestError(c, "validate failed") return } - if !mgr.requireImageOwner(c, req.ImageID) { - return - } for _, id := range req.IDList { if req.Type == "user" { if err := mgr.createImageUserEntity(c, req.ImageID, id); err != nil { @@ -742,9 +713,6 @@ func (mgr *ImagePackMgr) UserCancelShareImage(c *gin.Context) { resputil.BadRequestError(c, "validate failed") return } - if !mgr.requireImageOwner(c, req.ImageID) { - return - } if req.Type == "user" { if err := mgr.cancelShareImageWithUser(c, req.ImageID, req.ID); err != nil { resputil.Error(c, fmt.Sprintf("%v", err), resputil.NotSpecified) @@ -762,9 +730,9 @@ func (mgr *ImagePackMgr) UserCancelShareImage(c *gin.Context) { //nolint:dupl // ignore duplicate code func (mgr *ImagePackMgr) cancelShareImageWithAccount(c *gin.Context, imageID, accountID uint) error { accountImageQuery := query.ImageAccount - accountQuery := query.Account + userQuery := query.User // check if the account exists - if _, err := accountQuery.WithContext(c).Where(accountQuery.ID.Eq(accountID)).First(); err != nil { + if _, err := userQuery.WithContext(c).Where(userQuery.ID.Eq(accountID)).First(); err != nil { return fmt.Errorf("account does not exist: %w", err) } // check if the image has been shared to this account @@ -821,9 +789,6 @@ func (mgr *ImagePackMgr) GetImageGrantedUserOrAccount(c *gin.Context) { resputil.BadRequestError(c, "validate failed") return } - if !mgr.requireImageOwner(c, req.ImageID) { - return - } grangtedAccounts := []ImageGrantedAccounts{} imageAccountQuery := query.ImageAccount @@ -880,9 +845,6 @@ func (mgr *ImagePackMgr) UserGetImageUngrantedAccounts(c *gin.Context) { resputil.BadRequestError(c, "validate failed") return } - if !mgr.requireImageOwner(c, req.ImageID) { - return - } // 1. 查询已分享的AccountID sharedAccountIDs := []uint{} imageAccountQuery := query.ImageAccount @@ -943,9 +905,6 @@ func (mgr *ImagePackMgr) UserSearchUngrantedUsers(c *gin.Context) { resputil.BadRequestError(c, "validate failed") return } - if !mgr.requireImageOwner(c, req.ImageID) { - return - } // 1. 查询已分享的用户ID sharedUserIDs := []uint{} @@ -1018,7 +977,7 @@ func (mgr *ImagePackMgr) UserGetCudaBaseImages(c *gin.Context) { resputil.Success(c, resp) } -// AdminAddCudaBaseImage godoc +// UserAddCudaBaseImage godoc // // @Summary 添加CUDA基础镜像 // @Description 添加新的CUDA基础镜像到系统中 @@ -1027,8 +986,8 @@ func (mgr *ImagePackMgr) UserGetCudaBaseImages(c *gin.Context) { // @Produce json // @Security Bearer // @Param data body CudaBaseImageCreateRequest true "CUDA基础镜像信息" -// @Router /v1/admin/images/cudabaseimage [POST] -func (mgr *ImagePackMgr) AdminAddCudaBaseImage(c *gin.Context) { +// @Router /v1/images/cudabaseimage [POST] +func (mgr *ImagePackMgr) UserAddCudaBaseImage(c *gin.Context) { req := &CudaBaseImageCreateRequest{} if err := c.ShouldBindJSON(req); err != nil { klog.Errorf("validate cuda base image create request failed, err %v", err) @@ -1064,7 +1023,7 @@ func (mgr *ImagePackMgr) AdminAddCudaBaseImage(c *gin.Context) { resputil.Success(c, "cuda base image created successfully") } -// AdminDeleteCudaBaseImage godoc +// UserDeleteCudaBaseImage godoc // // @Summary 删除CUDA基础镜像 // @Description 根据ID删除指定的CUDA基础镜像 @@ -1073,8 +1032,8 @@ func (mgr *ImagePackMgr) AdminAddCudaBaseImage(c *gin.Context) { // @Produce json // @Security Bearer // @Param id path uint true "CUDA基础镜像ID" -// @Router /v1/admin/images/cudabaseimage/{id} [DELETE] -func (mgr *ImagePackMgr) AdminDeleteCudaBaseImage(c *gin.Context) { +// @Router /v1/images/cudabaseimage/{id} [DELETE] +func (mgr *ImagePackMgr) UserDeleteCudaBaseImage(c *gin.Context) { req := &CudaBaseImageDeleteRequest{} if err := c.ShouldBindUri(req); err != nil { klog.Errorf("validate cuda base image delete request failed, err %v", err) @@ -1123,9 +1082,6 @@ func (mgr *ImagePackMgr) UserUpdateImageArch(c *gin.Context) { resputil.BadRequestError(c, "validate failed") return } - if !mgr.requireImageOwner(c, req.ID) { - return - } mgr.updateImageArch(c, false, req.ID, req.Archs) } @@ -1173,21 +1129,3 @@ func (mgr *ImagePackMgr) updateImageArch(c *gin.Context, isAdminMode bool, image klog.Infof("image archs updated successfully, id: %d, new archs: %v", imageID, newArchs) resputil.Success(c, "image archs updated successfully") } - -func (mgr *ImagePackMgr) requireImageOwner(c *gin.Context, imageID uint) bool { - imageQuery := query.Image - if _, err := imageQuery.WithContext(c). - Where(imageQuery.ID.Eq(imageID)). - Where(imageQuery.UserID.Eq(util.GetToken(c).UserID)). - First(); err != nil { - if errors.Is(err, gorm.ErrRecordNotFound) { - klog.Warningf("permission denied or image not found, imageID: %d", imageID) - resputil.HandleError(c, bizerr.Forbidden.PermissionDenied.New("permission denied to manage this image")) - return false - } - klog.Errorf("check image owner failed, imageID: %d, err: %v", imageID, err) - resputil.HandleError(c, bizerr.Internal.DatabaseError.Wrap(err, "check image owner failed")) - return false - } - return true -} diff --git a/backend/internal/handler/image/interface.go b/backend/internal/handler/image/interface.go index 162b85afc..7c0518f78 100644 --- a/backend/internal/handler/image/interface.go +++ b/backend/internal/handler/image/interface.go @@ -63,6 +63,8 @@ func (mgr *ImagePackMgr) RegisterProtected(g *gin.RouterGroup) { g.GET("/user", mgr.UserSearchUngrantedUsers) g.GET("/account", mgr.UserGetImageUngrantedAccounts) g.GET("/cudabaseimage", mgr.UserGetCudaBaseImages) + g.POST("/cudabaseimage", mgr.UserAddCudaBaseImage) + g.DELETE("/cudabaseimage/:id", mgr.UserDeleteCudaBaseImage) g.POST("/arch", mgr.UserUpdateImageArch) } @@ -76,8 +78,6 @@ func (mgr *ImagePackMgr) RegisterAdmin(g *gin.RouterGroup) { g.POST("/description", mgr.AdminChangeImageDescription) g.POST("/tags", mgr.AdminChangeImageTags) g.POST("/arch", mgr.AdminUpdateImageArch) - g.POST("/cudabaseimage", mgr.AdminAddCudaBaseImage) - g.DELETE("/cudabaseimage/:id", mgr.AdminDeleteCudaBaseImage) } func NewImagePackMgr(conf *handler.RegisterConfig) handler.Manager { diff --git a/backend/internal/handler/image/interface_test.go b/backend/internal/handler/image/interface_test.go deleted file mode 100644 index 62ef8e774..000000000 --- a/backend/internal/handler/image/interface_test.go +++ /dev/null @@ -1,42 +0,0 @@ -package image - -import ( - "net/http" - "testing" - - "github.com/gin-gonic/gin" -) - -func TestCudaBaseImageWriteRoutesAreAdminOnly(t *testing.T) { - gin.SetMode(gin.TestMode) - router := gin.New() - mgr := &ImagePackMgr{} - mgr.RegisterProtected(router.Group("/api/v1/images")) - mgr.RegisterAdmin(router.Group("/api/v1/admin/images")) - - routes := make(map[string]bool) - for _, route := range router.Routes() { - routes[route.Method+" "+route.Path] = true - } - - want := []string{ - http.MethodGet + " /api/v1/images/cudabaseimage", - http.MethodPost + " /api/v1/admin/images/cudabaseimage", - http.MethodDelete + " /api/v1/admin/images/cudabaseimage/:id", - } - for _, route := range want { - if !routes[route] { - t.Errorf("missing route %s", route) - } - } - - forbidden := []string{ - http.MethodPost + " /api/v1/images/cudabaseimage", - http.MethodDelete + " /api/v1/images/cudabaseimage/:id", - } - for _, route := range forbidden { - if routes[route] { - t.Errorf("user route must not expose CUDA mutation: %s", route) - } - } -} diff --git a/backend/internal/handler/interface.go b/backend/internal/handler/interface.go index 91016b8cb..7c1dff9b0 100644 --- a/backend/internal/handler/interface.go +++ b/backend/internal/handler/interface.go @@ -1,12 +1,20 @@ package handler import ( + "context" + "encoding/json" + "time" + "github.com/gin-gonic/gin" + corev1 "k8s.io/api/core/v1" "k8s.io/client-go/kubernetes" "k8s.io/client-go/rest" "sigs.k8s.io/controller-runtime/pkg/client" + batch "volcano.sh/apis/pkg/apis/batch/v1alpha1" + "github.com/raids-lab/crater/dao/model" "github.com/raids-lab/crater/internal/service" + "github.com/raids-lab/crater/internal/util" "github.com/raids-lab/crater/pkg/aitaskctl" "github.com/raids-lab/crater/pkg/crclient" "github.com/raids-lab/crater/pkg/cronjob" @@ -16,6 +24,33 @@ import ( "github.com/raids-lab/crater/pkg/prequeuewatcher" ) +type JobContextResp struct { + Meta struct { + Name string `json:"name,omitempty"` + JobName string `json:"jobName,omitempty"` + Namespace string `json:"namespace,omitempty"` + User string `json:"user,omitempty"` + Queue string `json:"queue,omitempty"` + JobType model.JobType `json:"jobType,omitempty"` + Status batch.JobPhase `json:"status,omitempty"` + CreationTimestamp time.Time `json:"creationTimestamp,omitempty"` + RunningTimestamp time.Time `json:"runningTimestamp,omitempty"` + CompletedTimestamp time.Time `json:"completedTimestamp,omitempty"` + Nodes []string `json:"nodes,omitempty"` + Resources corev1.ResourceList `json:"resources,omitempty"` + } `json:"meta"` + DB struct { + ProfileData any `json:"profileData,omitempty"` + ScheduleData any `json:"scheduleData,omitempty"` + Events []corev1.Event `json:"events,omitempty"` + TerminatedStates any `json:"terminatedStates,omitempty"` + } `json:"db"` + Log struct { + Container string `json:"container,omitempty"` + Tail string `json:"tail,omitempty"` + } `json:"log,omitempty"` +} + // Manager is the interface that wraps the basic methods for a handler manager. type Manager interface { GetName() string @@ -60,6 +95,84 @@ type RegisterConfig struct { GpuAnalysisService *service.GpuAnalysisService } +type JobMutationSubmitter interface { + SubmitJupyterJob(ctx context.Context, token util.JWTMessage, req json.RawMessage) (any, error) + SubmitWebIDEJob(ctx context.Context, token util.JWTMessage, req json.RawMessage) (any, error) + SubmitTrainingJob(ctx context.Context, token util.JWTMessage, req json.RawMessage) (any, error) + SubmitPytorchJob(ctx context.Context, token util.JWTMessage, req json.RawMessage) (any, error) + SubmitTensorflowJob(ctx context.Context, token util.JWTMessage, req json.RawMessage) (any, error) + DeleteJob(ctx context.Context, token util.JWTMessage, jobName string) (any, error) + StopJob(ctx context.Context, token util.JWTMessage, jobName string) (any, error) + ResubmitJob(ctx context.Context, token util.JWTMessage, req json.RawMessage) (any, error) +} + +type JobInsightReader interface { + FindScopedJob(ctx context.Context, token util.JWTMessage, jobName string) (*model.Job, error) + BuildJobDetail(job *model.Job) any + GetJobEvents(ctx context.Context, token util.JWTMessage, jobName string) (any, error) + GetJobLog(ctx context.Context, token util.JWTMessage, jobName string, tailLines int64, keyword string) (map[string]string, error) + GetDiagnosticContext( + ctx context.Context, + token util.JWTMessage, + jobName string, + includeLog bool, + tailLines int64, + ) (JobContextResp, error) +} + +type ImageAccessRecord struct { + Image *model.Image + ShareStatus model.ImageShareType +} + +type ImageInsightReader interface { + ListAccessibleImages(ctx context.Context, token util.JWTMessage) ([]ImageAccessRecord, error) +} + +var jobMutationSubmitterFactory func(conf *RegisterConfig) JobMutationSubmitter +var jobInsightReaderFactory func(conf *RegisterConfig) JobInsightReader +var imageInsightReaderFactory func(conf *RegisterConfig) ImageInsightReader + +func RegisterJobMutationSubmitterFactory(factory func(conf *RegisterConfig) JobMutationSubmitter) { + jobMutationSubmitterFactory = factory +} + +func RegisterJobInsightReaderFactory(factory func(conf *RegisterConfig) JobInsightReader) { + jobInsightReaderFactory = factory +} + +func RegisterImageInsightReaderFactory(factory func(conf *RegisterConfig) ImageInsightReader) { + imageInsightReaderFactory = factory +} + +func NewJobMutationSubmitter(conf *RegisterConfig) JobMutationSubmitter { + if jobMutationSubmitterFactory == nil { + return nil + } + return jobMutationSubmitterFactory(conf) +} + +func NewJobInsightReader(conf *RegisterConfig) JobInsightReader { + if jobInsightReaderFactory == nil { + return nil + } + return jobInsightReaderFactory(conf) +} + +func NewImageInsightReader(conf *RegisterConfig) ImageInsightReader { + if imageInsightReaderFactory == nil { + return nil + } + return imageInsightReaderFactory(conf) +} + +// InternalRouter is an optional interface for managers that expose internal-only endpoints +// (e.g. service-to-service callbacks authenticated via X-Agent-Internal-Token). +// Managers that do not need internal routes do not need to implement this interface. +type InternalRouter interface { + RegisterInternal(group *gin.RouterGroup) +} + // Registers is a slice of Manager Init functions. // Each Manager should register itself by appending its Init function to this slice. var Registers = []func(config *RegisterConfig) Manager{} diff --git a/backend/internal/handler/modeldownload.go b/backend/internal/handler/modeldownload.go index dee2d3759..35fed2b0d 100644 --- a/backend/internal/handler/modeldownload.go +++ b/backend/internal/handler/modeldownload.go @@ -1,7 +1,7 @@ package handler import ( - "context" + "crypto/sha256" "errors" "fmt" "io" @@ -30,7 +30,6 @@ import ( "github.com/raids-lab/crater/dao/query" "github.com/raids-lab/crater/internal/bizerr" "github.com/raids-lab/crater/internal/resputil" - "github.com/raids-lab/crater/internal/service" "github.com/raids-lab/crater/internal/util" "github.com/raids-lab/crater/pkg/config" "github.com/raids-lab/crater/pkg/utils" @@ -50,18 +49,16 @@ func init() { } type ModelDownloadMgr struct { - name string - crClient kubernetes.Interface - namespace string - quotaService *service.ModelDownloadQuotaService + name string + crClient kubernetes.Interface + namespace string } func NewModelDownloadMgr(conf *RegisterConfig) Manager { return &ModelDownloadMgr{ - name: "model-download", - crClient: conf.KubeClient, - namespace: config.GetConfig().Namespaces.Job, - quotaService: service.NewModelDownloadQuotaService(conf.ConfigService), + name: "model-download", + crClient: conf.KubeClient, + namespace: config.GetConfig().Namespaces.Job, } } @@ -99,52 +96,39 @@ type CreateDownloadReq struct { type DownloadActionReq struct { Token string `json:"token"` - // Revision is optional and only used by retry. A non-nil empty value means - // "use the source default branch" while preserving the failed record's path. - Revision *string `json:"revision"` } type ModelDownloadResp struct { - ID uint `json:"id"` - Name string `json:"name"` - Source string `json:"source"` - Category string `json:"category"` - Revision string `json:"revision"` - Path string `json:"path"` - SizeBytes int64 `json:"sizeBytes"` - DownloadedBytes int64 `json:"downloadedBytes"` - DownloadSpeed string `json:"downloadSpeed"` - Status string `json:"status"` - Message string `json:"message"` - JobName string `json:"jobName"` - CreatorID uint `json:"creatorId"` - ReferenceCount int `json:"referenceCount"` // Deprecated: use requesterCount. - RequesterCount int `json:"requesterCount"` - Requesters []model.UserInfo `json:"requesters"` - Relation string `json:"relation"` - CreatedAt time.Time `json:"createdAt"` - UpdatedAt time.Time `json:"updatedAt"` - SourceUpdatedAt *time.Time `json:"sourceUpdatedAt"` - UserInfo model.UserInfo `json:"userInfo"` - CanManage bool `json:"canManage"` - CanDelete bool `json:"canDelete"` - CanViewLogs bool `json:"canViewLogs"` - SourceURL string `json:"sourceUrl"` - DisplayName string `json:"displayName"` - License string `json:"license"` - Task string `json:"task"` - Library string `json:"library"` - ModelType string `json:"modelType"` - ParameterCount int64 `json:"parameterCount"` - SourceCreatedAt *time.Time `json:"sourceCreatedAt"` + ID uint `json:"id"` + Name string `json:"name"` + Source string `json:"source"` + Category string `json:"category"` + Revision string `json:"revision"` + Path string `json:"path"` + SizeBytes int64 `json:"sizeBytes"` + DownloadedBytes int64 `json:"downloadedBytes"` + DownloadSpeed string `json:"downloadSpeed"` + Status string `json:"status"` + Message string `json:"message"` + JobName string `json:"jobName"` + CreatorID uint `json:"creatorId"` + ReferenceCount int `json:"referenceCount"` + CreatedAt time.Time `json:"createdAt"` + UpdatedAt time.Time `json:"updatedAt"` + SourceUpdatedAt *time.Time `json:"sourceUpdatedAt"` + UserInfo model.UserInfo `json:"userInfo"` + CanManage bool `json:"canManage"` + CanViewLogs bool `json:"canViewLogs"` + SourceURL string `json:"sourceUrl"` + DisplayName string `json:"displayName"` + License string `json:"license"` + Task string `json:"task"` + Library string `json:"library"` + ModelType string `json:"modelType"` + ParameterCount int64 `json:"parameterCount"` + SourceCreatedAt *time.Time `json:"sourceCreatedAt"` } -const ( - ModelDownloadRelationNone = "none" - ModelDownloadRelationCreator = "creator" - ModelDownloadRelationSubmitted = "submitted" -) - type ListDownloadsReq struct { Page int `form:"page"` PageSize int `form:"pageSize,default=20"` @@ -192,28 +176,21 @@ func (mgr *ModelDownloadMgr) canViewDownloadLogs( return false, bizerr.Internal.DatabaseError.Wrap(err, "check download log permission failed") } -func (mgr *ModelDownloadMgr) applyDownloadUserContext( +func (mgr *ModelDownloadMgr) applyLogViewPermissions( c *gin.Context, downloads []*model.ModelDownload, responses []ModelDownloadResp, token util.JWTMessage, ) error { - if len(downloads) != len(responses) { - return bizerr.Internal.DatabaseError.New("download responses do not match download records") - } - downloadsByID := make(map[uint]*model.ModelDownload, len(downloads)) - for _, download := range downloads { - downloadsByID[download.ID] = download + if token.RolePlatform == model.RoleAdmin { + for i := range responses { + responses[i].CanViewLogs = true + } + return nil } + ids := make([]uint, 0, len(downloads)) - for i := range responses { - response := &responses[i] - download, ok := downloadsByID[response.ID] - if !ok { - return bizerr.Internal.DatabaseError.New("download response does not match download record") - } - if token.RolePlatform == model.RoleAdmin { - response.CanViewLogs = true - } + for i, download := range downloads { if download.CreatorID == token.UserID { - response.CanViewLogs = true + responses[i].CanViewLogs = true + continue } ids = append(ids, download.ID) } @@ -223,53 +200,22 @@ func (mgr *ModelDownloadMgr) applyDownloadUserContext( q := query.UserModelDownload associations, err := q.WithContext(c). - Preload(q.User). - Where(q.ModelDownloadID.In(ids...)). - Order(q.CreatedAt.Asc()). + Where(q.UserID.Eq(token.UserID), q.ModelDownloadID.In(ids...)). Find() if err != nil { - return bizerr.Internal.DatabaseError.Wrap(err, "list download requesters failed") + return bizerr.Internal.DatabaseError.Wrap(err, "list download log permissions failed") } - applyDownloadRequesters(responses, associations, token) - return nil -} - -func applyDownloadRequesters( - responses []ModelDownloadResp, associations []*model.UserModelDownload, token util.JWTMessage, -) { - requestersByDownload := make(map[uint][]model.UserInfo, len(responses)) - requesterCountByDownload := make(map[uint]int, len(responses)) - requestedByCurrentUser := make(map[uint]struct{}, len(responses)) + allowed := make(map[uint]struct{}, len(associations)) for _, association := range associations { - requesterCountByDownload[association.ModelDownloadID]++ - if association.UserID == token.UserID { - requestedByCurrentUser[association.ModelDownloadID] = struct{}{} - } - if association.User.Name == "" { - continue - } - nickname := association.User.Nickname - if nickname == "" { - nickname = association.User.Name - } - requestersByDownload[association.ModelDownloadID] = append( - requestersByDownload[association.ModelDownloadID], - model.UserInfo{Username: association.User.Name, Nickname: nickname}, - ) + allowed[association.ModelDownloadID] = struct{}{} } - for i := range responses { - response := &responses[i] - response.RequesterCount = requesterCountByDownload[response.ID] - response.ReferenceCount = response.RequesterCount - response.Requesters = requestersByDownload[response.ID] - if response.Requesters == nil { - response.Requesters = []model.UserInfo{} - } - if _, ok := requestedByCurrentUser[response.ID]; ok && response.Relation != ModelDownloadRelationCreator { - response.Relation = ModelDownloadRelationSubmitted - response.CanViewLogs = true + for i, download := range downloads { + if responses[i].CanViewLogs { + continue } + _, responses[i].CanViewLogs = allowed[download.ID] } + return nil } // CreateDownload godoc @@ -293,99 +239,28 @@ func (mgr *ModelDownloadMgr) associateUserWithDownload( return nil } -// findReadyOrOngoingDownload only reuses a download with the exact requested -// upstream identity. The canonical storage path is shared by all variants, so -// a different source or revision must be reported as a conflict instead of -// silently satisfying the request with unrelated files. +// findReadyOrOngoingDownload 查找已完成或正在进行的下载 func (mgr *ModelDownloadMgr) findReadyOrOngoingDownload( c *gin.Context, txQ *query.Query, name string, source model.ModelSource, category model.DownloadCategory, revision string, -) (*model.ModelDownload, error) { +) *model.ModelDownload { q := txQ.ModelDownload - readyDownload, err := q.WithContext(c). - Clauses(clause.Locking{Strength: "UPDATE"}). - Where(q.Name.Eq(name), q.Source.Eq(string(source)), q.Category.Eq(string(category)), - q.Revision.Eq(revision), - q.Status.Eq(string(model.ModelDownloadStatusReady))). - First() - if err == nil { - return readyDownload, nil - } - if !errors.Is(err, gorm.ErrRecordNotFound) { - return nil, bizerr.Internal.DatabaseError.Wrap(err, "find ready logical download") + // 查询下载成功的记录 + existingDownload, _ := q.WithContext(c). + Where(q.Name.Eq(name), q.Source.Eq(string(source)), + q.Category.Eq(string(category)), q.Revision.Eq(revision), + q.Status.Eq(string(model.ModelDownloadStatusReady))).First() + if existingDownload != nil { + return existingDownload } - ongoingDownload, err := q.WithContext(c). - Clauses(clause.Locking{Strength: "UPDATE"}). + // 查询正在进行的下载 + ongoingDownload, _ := q.WithContext(c). Where(q.Name.Eq(name), q.Source.Eq(string(source)), q.Category.Eq(string(category)), - q.Revision.Eq(revision), - q.Status.In(string(model.ModelDownloadStatusPending), string(model.ModelDownloadStatusDownloading), - string(model.ModelDownloadStatusPaused))). - First() - if err == nil { - return ongoingDownload, nil - } - if !errors.Is(err, gorm.ErrRecordNotFound) { - return nil, bizerr.Internal.DatabaseError.Wrap(err, "find ongoing logical download") - } - return nil, nil -} - -// lockModelDownloadIdentity serializes first-time downloads whose current -// database uniqueness still differs by source and revision. Existing -// installations may contain more than one historical record, so changing the -// unique index in place would make upgrades fail; a PostgreSQL transaction -// advisory lock prevents two new variants from racing into the same short path. -func lockModelDownloadIdentity( - c *gin.Context, txQ *query.Query, name string, category model.DownloadCategory, -) error { - db := txQ.ModelDownload.WithContext(c).UnderlyingDB() - if db.Name() != "postgres" { - return nil - } - identity := string(category) + ":" + name - if err := db.Exec("SELECT pg_advisory_xact_lock(hashtextextended(?, 0))", identity).Error; err != nil { - return bizerr.Internal.DatabaseError.Wrap(err, "lock logical download identity") - } - return nil -} - -func updateDownloadAndReleaseQuota( - ctx context.Context, downloadID uint, updates map[string]any, -) error { - return query.GetDB().WithContext(ctx).Transaction(func(tx *gorm.DB) error { - if err := tx.Model(&model.ModelDownload{}).Where("id = ?", downloadID).Updates(updates).Error; err != nil { - return err - } - return service.ReleaseModelDownloadQuotaReservation(ctx, tx, downloadID) - }) -} - -// checkLogicalDownloadConflict blocks a new source/revision when a historical -// failed or soft-deleted record already owns storage for the same public model. -// Ready and ongoing records are handled earlier and reused instead. -func checkLogicalDownloadConflict( - c *gin.Context, txQ *query.Query, - name string, source model.ModelSource, category model.DownloadCategory, revision string, -) error { - var conflict model.ModelDownload - db := txQ.ModelDownload.WithContext(c).Unscoped().UnderlyingDB() - err := db.Clauses(clause.Locking{Strength: "UPDATE"}). - Where("name = ? AND category = ? AND NOT (source = ? AND revision = ?)", - name, category, source, revision). - Order("id ASC"). - First(&conflict).Error - if err == nil { - return bizerr.Conflict.ResourceStatusError.New(fmt.Sprintf( - "model already has storage at %s from %s revision %q; reuse or resolve that record before downloading another source or revision", - conflict.Path, conflict.Source, conflict.Revision, - )) - } - if !errors.Is(err, gorm.ErrRecordNotFound) { - return bizerr.Internal.DatabaseError.Wrap(err, "check logical download conflict") - } - return nil + q.Revision.Eq(revision), q.Status.In(string(model.ModelDownloadStatusPending), + string(model.ModelDownloadStatusDownloading))).First() + return ongoingDownload } // restoreAndResetSoftDeletedDownload 恢复并重置软删除的下载记录 @@ -476,17 +351,8 @@ func (mgr *ModelDownloadMgr) getOrCreateDownload( var isNewDownload bool err := db.Transaction(func(tx *query.Query) error { - if err := lockModelDownloadIdentity(c, tx, req.Name, category); err != nil { - return err - } - // 1. 查找已完成或正在进行的下载 - existing, err := mgr.findReadyOrOngoingDownload( - c, tx, req.Name, source, category, revision, - ) - if err != nil { - return err - } + existing := mgr.findReadyOrOngoingDownload(c, tx, req.Name, source, category, revision) if existing != nil { if err := mgr.associateUserWithDownload(c, tx, token.UserID, existing.ID); err != nil { return err @@ -495,10 +361,6 @@ func (mgr *ModelDownloadMgr) getOrCreateDownload( return nil } - if err := checkLogicalDownloadConflict(c, tx, req.Name, source, category, revision); err != nil { - return err - } - // 2. 查找并恢复软删除的记录 restored, err := mgr.restoreAndResetSoftDeletedDownload(c, tx, req.Name, source, category, revision, token.Username) if err != nil { @@ -508,13 +370,7 @@ func (mgr *ModelDownloadMgr) getOrCreateDownload( if err := mgr.associateUserWithDownload(c, tx, token.UserID, restored.ID); err != nil { return err } - download, isNewDownload = restored, shouldSubmitRestoredDownload(restored) - if isNewDownload { - return mgr.quotaService.Reserve( - c, tx.ModelDownload.WithContext(c).UnderlyingDB(), token.UserID, - restored.ID, model.ModelDownloadSubmissionCreate, - ) - } + download, isNewDownload = restored, true return nil } @@ -528,10 +384,7 @@ func (mgr *ModelDownloadMgr) getOrCreateDownload( return err } download, isNewDownload = failed, true - return mgr.quotaService.Reserve( - c, tx.ModelDownload.WithContext(c).UnderlyingDB(), token.UserID, - failed.ID, model.ModelDownloadSubmissionCreate, - ) + return nil } // 4. 创建新的下载任务 @@ -550,19 +403,12 @@ func (mgr *ModelDownloadMgr) getOrCreateDownload( return fmt.Errorf("create association failed: %w", err) } download, isNewDownload = newDownload, true - return mgr.quotaService.Reserve( - c, tx.ModelDownload.WithContext(c).UnderlyingDB(), token.UserID, - newDownload.ID, model.ModelDownloadSubmissionCreate, - ) + return nil }) return download, isNewDownload, err } -func shouldSubmitRestoredDownload(download *model.ModelDownload) bool { - return download.Status != model.ModelDownloadStatusReady -} - // @Summary 创建模型下载任务 // @Description 创建一个新的模型下载任务 // @Tags ModelDownload @@ -613,13 +459,13 @@ func (mgr *ModelDownloadMgr) CreateDownload(c *gin.Context) { } else { basePath = "public/Datasets" } - downloadPath := modelDownloadStoragePath(basePath, safeName) + downloadPath := modelDownloadStoragePath(basePath, safeName, source, req.Revision) // 在事务中获取或创建下载任务 download, isNewDownload, err := mgr.getOrCreateDownload(c, &req, token, source, category, downloadPath, req.Revision) if err != nil { klog.Errorf("get or create download failed: %v", err) - resputil.HandleError(c, err) + resputil.HandleError(c, bizerr.Internal.ServiceError.Wrap(err, "处理下载请求失败")) return } @@ -627,9 +473,6 @@ func (mgr *ModelDownloadMgr) CreateDownload(c *gin.Context) { if !isNewDownload { resp := convertDownloadToResp(download, token) resp.CanViewLogs = true - if download.CreatorID != token.UserID { - resp.Relation = ModelDownloadRelationSubmitted - } if download.Status == model.ModelDownloadStatusReady { c.JSON(http.StatusOK, gin.H{ "code": resputil.OK, @@ -640,7 +483,7 @@ func (mgr *ModelDownloadMgr) CreateDownload(c *gin.Context) { c.JSON(http.StatusOK, gin.H{ "code": resputil.OK, "data": resp, - "msg": "该资源正在下载中，已记录您的下载需求", + "msg": "该资源正在下载中，已将您加入共享列表", }) } return @@ -649,13 +492,12 @@ func (mgr *ModelDownloadMgr) CreateDownload(c *gin.Context) { // 提交 K8s Job if err := mgr.submitDownloadJob(c, download, token.Username, req.Token); err != nil { klog.Errorf("submit download job failed: %v", err) + q := query.ModelDownload updates := map[string]any{ "status": model.ModelDownloadStatusFailed, "message": fmt.Sprintf("submit job failed: %v", err), } - if rollbackErr := updateDownloadAndReleaseQuota(c, download.ID, updates); rollbackErr != nil { - klog.Errorf("release failed download quota reservation: %v", rollbackErr) - } + _, _ = q.WithContext(c).Where(q.ID.Eq(download.ID)).Updates(updates) resputil.Error(c, "submit download job failed", resputil.NotSpecified) return @@ -711,7 +553,7 @@ func (mgr *ModelDownloadMgr) ListDownloads(c *gin.Context) { for i, d := range downloads { resp[i] = convertDownloadToResp(d, token) } - if err := mgr.applyDownloadUserContext(c, downloads, resp, token); err != nil { + if err := mgr.applyLogViewPermissions(c, downloads, resp, token); err != nil { resputil.HandleError(c, err) return } @@ -752,7 +594,7 @@ func (mgr *ModelDownloadMgr) ListDownloads(c *gin.Context) { for i, d := range downloads { items[i] = convertDownloadToResp(d, token) } - if err := mgr.applyDownloadUserContext(c, downloads, items, token); err != nil { + if err := mgr.applyLogViewPermissions(c, downloads, items, token); err != nil { resputil.HandleError(c, err) return } @@ -817,12 +659,12 @@ func (mgr *ModelDownloadMgr) GetDownload(c *gin.Context) { } resp := convertDownloadToResp(download, token) - responses := []ModelDownloadResp{resp} - if err := mgr.applyDownloadUserContext(c, []*model.ModelDownload{download}, responses, token); err != nil { + canViewLogs, err := mgr.canViewDownloadLogs(c, download.ID, token) + if err != nil { resputil.HandleError(c, err) return } - resp = responses[0] + resp.CanViewLogs = canViewLogs resputil.Success(c, resp) } @@ -852,7 +694,7 @@ func (mgr *ModelDownloadMgr) requireManagePermission(c *gin.Context, downloadID // @Produce json // @Security Bearer // @Param id path int true "下载任务ID" -// @Param data body DownloadActionReq false "可选的临时访问令牌和重试版本" +// @Param data body DownloadActionReq false "可选的临时访问令牌" // @Success 200 {object} resputil.Response[ModelDownloadResp] // @Router /v1/models/downloads/{id}/retry [POST] func (mgr *ModelDownloadMgr) RetryDownload(c *gin.Context) { @@ -870,149 +712,86 @@ func (mgr *ModelDownloadMgr) RetryDownload(c *gin.Context) { resputil.HandleError(c, err) return } - if err := normalizeRetryRevision(&action); err != nil { - resputil.HandleError(c, err) - return - } if _, err := mgr.requireManagePermission(c, req.ID); err != nil { resputil.HandleError(c, err) return } - download, err := mgr.prepareRetryDownload(c, req.ID, token, action.Revision) - - if err != nil { - klog.Errorf("retry download transaction failed: %v", err) - resputil.HandleError(c, err) - return - } - // 事务成功后提交 Job - if err := mgr.submitDownloadJob(c, download, token.Username, action.Token); err != nil { - klog.Errorf("submit download job failed: %v", err) - // 回滚状态 - updates := map[string]any{ - "status": model.ModelDownloadStatusFailed, - "message": fmt.Sprintf("submit job failed: %v", err), - } - if rollbackErr := updateDownloadAndReleaseQuota(c, download.ID, updates); rollbackErr != nil { - klog.Errorf("release failed retry quota reservation: %v", rollbackErr) - } - resputil.HandleError(c, bizerr.Internal.K8sServiceError.Wrap(err, "submit download job failed")) - return - } - - download.Status = model.ModelDownloadStatusDownloading - download.Message = "" - - resputil.Success(c, convertDownloadToResp(download, token)) -} - -func normalizeRetryRevision(action *DownloadActionReq) error { - if action.Revision == nil { - return nil - } - revision := strings.TrimSpace(*action.Revision) - if len(revision) > maxDownloadRevisionLength { - return bizerr.BadRequest.ParameterError.New("revision must not exceed 128 characters") - } - action.Revision = &revision - return nil -} + q := query.ModelDownload -// prepareRetryDownload updates the failed record under a row lock. Correcting -// the revision deliberately does not rewrite Path: the source SDK receives the -// same local directory and can validate/reuse files from the failed attempt. -func (mgr *ModelDownloadMgr) prepareRetryDownload( - c *gin.Context, downloadID uint, token util.JWTMessage, revision *string, -) (*model.ModelDownload, error) { + // 使用事务和行锁来防止并发重试导致的竞态条件 db := query.Use(query.GetDB()) var download *model.ModelDownload - err := db.Transaction(func(tx *query.Query) error { + var newJobName string + + err = db.Transaction(func(tx *query.Query) error { txQ := tx.ModelDownload + + // 使用 FOR UPDATE 锁定记录 d, err := txQ.WithContext(c). Clauses(clause.Locking{Strength: "UPDATE"}). - Where(txQ.ID.Eq(downloadID)). + Where(txQ.ID.Eq(req.ID)). First() if err != nil { return bizerr.NotFound.DataBaseNotFound.Wrap(err, "download not found") } + + // 再次检查状态（可能被其他请求修改了） if d.Status != model.ModelDownloadStatusFailed { return bizerr.Conflict.ResourceStatusError.New( fmt.Sprintf("only failed downloads can be retried, current status: %s", d.Status), ) } - if err := checkRetryRevisionConflict(c, tx, d, revision); err != nil { - return err - } - if err := mgr.quotaService.Reserve( - c, tx.ModelDownload.WithContext(c).UnderlyingDB(), token.UserID, - d.ID, model.ModelDownloadSubmissionRetry, - ); err != nil { - return err - } - if err := mgr.associateUserWithDownload(c, tx, token.UserID, d.ID); err != nil { - return err - } - newJobName := fmt.Sprintf("model-dl-%s-%s", token.Username, uuid.New().String()[:8]) + // 生成新的 Job 名称并更新 + newJobName = fmt.Sprintf("model-dl-%s-%s", token.Username, uuid.New().String()[:8]) updates := map[string]any{ - "status": model.ModelDownloadStatusDownloading, "message": "", "job_name": newJobName, - } - if revision != nil { - updates["revision"] = *revision + "status": model.ModelDownloadStatusDownloading, + "message": "", + "job_name": newJobName, } - if _, err = txQ.WithContext(c).Where(txQ.ID.Eq(d.ID)).Updates(updates); err != nil { - return retryUpdateError(err) + _, err = txQ.WithContext(c).Where(txQ.ID.Eq(d.ID)).Updates(updates) + if err != nil { + return bizerr.Internal.DatabaseError.Wrap(err, "update download record failed") } - d.Status, d.JobName, d.Message = model.ModelDownloadStatusDownloading, newJobName, "" - if revision != nil { - d.Revision = *revision - } + d.Status = model.ModelDownloadStatusDownloading + d.JobName = newJobName + d.Message = "" download = d return nil }) - return download, err -} -func checkRetryRevisionConflict( - c *gin.Context, txQ *query.Query, download *model.ModelDownload, revision *string, -) error { - if revision == nil || *revision == download.Revision { - return nil - } - conflict, err := txQ.ModelDownload.WithContext(c).Unscoped(). - Clauses(clause.Locking{Strength: "UPDATE"}). - Where(txQ.ModelDownload.ID.Neq(download.ID), txQ.ModelDownload.Name.Eq(download.Name), - txQ.ModelDownload.Source.Eq(string(download.Source)), - txQ.ModelDownload.Category.Eq(string(download.Category)), - txQ.ModelDownload.Revision.Eq(*revision)). - First() - if err == nil && conflict != nil { - return bizerr.Conflict.ResourceStatusError.New("download with requested revision already exists") - } - if err != nil && !errors.Is(err, gorm.ErrRecordNotFound) { - return bizerr.Internal.DatabaseError.Wrap(err, "check retry revision conflict") + if err != nil { + klog.Errorf("retry download transaction failed: %v", err) + resputil.HandleError(c, err) + return } - return nil -} -func retryUpdateError(err error) error { - if errors.Is(err, gorm.ErrDuplicatedKey) { - return bizerr.Conflict.ResourceStatusError.New("download with requested revision already exists") - } - var sqlStateError interface{ SQLState() string } - if errors.As(err, &sqlStateError) && sqlStateError.SQLState() == "23505" { - return bizerr.Conflict.ResourceStatusError.New("download with requested revision already exists") + // 事务成功后提交 Job + if err := mgr.submitDownloadJob(c, download, token.Username, action.Token); err != nil { + klog.Errorf("submit download job failed: %v", err) + // 回滚状态 + updates := map[string]any{ + "status": model.ModelDownloadStatusFailed, + "message": fmt.Sprintf("submit job failed: %v", err), + } + _, _ = q.WithContext(c).Where(q.ID.Eq(download.ID)).Updates(updates) + resputil.HandleError(c, bizerr.Internal.K8sServiceError.Wrap(err, "submit download job failed")) + return } - return bizerr.Internal.DatabaseError.Wrap(err, "update download record failed") + + download.Status = model.ModelDownloadStatusDownloading + download.Message = "" + + resputil.Success(c, convertDownloadToResp(download, token)) } // DeleteDownload godoc // // @Summary 删除模型下载任务 -// @Description 删除下载任务记录(仅平台管理员),已下载的文件保留在存储中 +// @Description 删除下载任务记录(仅创建者或管理员),已下载的文件保留在存储中 // @Tags ModelDownload // @Accept json // @Produce json @@ -1021,11 +800,6 @@ func retryUpdateError(err error) error { // @Success 200 {object} resputil.Response[string] // @Router /v1/models/downloads/{id} [DELETE] func (mgr *ModelDownloadMgr) DeleteDownload(c *gin.Context) { - token := util.GetToken(c) - if !canDeleteDownload(token) { - resputil.HandleError(c, bizerr.Forbidden.PermissionDenied.New("only a platform administrator can delete download records")) - return - } var req struct { ID uint `uri:"id" binding:"required"` } @@ -1034,10 +808,9 @@ func (mgr *ModelDownloadMgr) DeleteDownload(c *gin.Context) { resputil.HandleError(c, bizerr.BadRequest.ParameterError.Wrap(err, "invalid download id")) return } - q := query.ModelDownload - download, err := q.WithContext(c).Where(q.ID.Eq(req.ID)).First() + download, err := mgr.requireManagePermission(c, req.ID) if err != nil { - resputil.HandleError(c, bizerr.NotFound.DataBaseNotFound.Wrap(err, "download not found")) + resputil.HandleError(c, err) return } @@ -1060,10 +833,6 @@ func (mgr *ModelDownloadMgr) DeleteDownload(c *gin.Context) { resputil.Success(c, "deleted successfully") } -func canDeleteDownload(token util.JWTMessage) bool { - return token.RolePlatform == model.RoleAdmin -} - // PauseDownload godoc // // @Summary 暂停下载任务 @@ -1091,6 +860,8 @@ func (mgr *ModelDownloadMgr) PauseDownload(c *gin.Context) { return } + q := query.ModelDownload + if download.Status != model.ModelDownloadStatusDownloading { resputil.HandleError(c, bizerr.Conflict.ResourceStatusError.New("only downloading tasks can be paused")) return @@ -1107,7 +878,7 @@ func (mgr *ModelDownloadMgr) PauseDownload(c *gin.Context) { "status": model.ModelDownloadStatusPaused, "message": "Download paused by user", } - if err := updateDownloadAndReleaseQuota(c, download.ID, updates); err != nil { + if _, err := q.WithContext(c).Where(q.ID.Eq(download.ID)).Updates(updates); err != nil { resputil.HandleError(c, bizerr.Internal.DatabaseError.Wrap(err, "update paused download failed")) return } @@ -1152,6 +923,7 @@ func (mgr *ModelDownloadMgr) ResumeDownload(c *gin.Context) { return } + q := query.ModelDownload db := query.Use(query.GetDB()) err = db.Transaction(func(tx *query.Query) error { txQ := tx.ModelDownload @@ -1165,12 +937,6 @@ func (mgr *ModelDownloadMgr) ResumeDownload(c *gin.Context) { if locked.Status != model.ModelDownloadStatusPaused { return bizerr.Conflict.ResourceStatusError.New("only paused tasks can be resumed") } - if limitErr := mgr.quotaService.Reserve( - c, tx.ModelDownload.WithContext(c).UnderlyingDB(), token.UserID, - locked.ID, model.ModelDownloadSubmissionResume, - ); limitErr != nil { - return limitErr - } newJobName := fmt.Sprintf("model-dl-%s-%s", token.Username, uuid.New().String()[:8]) if _, updateErr := txQ.WithContext(c).Where(txQ.ID.Eq(locked.ID)).Updates(map[string]any{ @@ -1199,9 +965,7 @@ func (mgr *ModelDownloadMgr) ResumeDownload(c *gin.Context) { "status": model.ModelDownloadStatusPaused, "message": fmt.Sprintf("resume failed: %v", err), } - if rollbackErr := updateDownloadAndReleaseQuota(c, download.ID, rollbackUpdates); rollbackErr != nil { - klog.Errorf("release failed resume quota reservation: %v", rollbackErr) - } + _, _ = q.WithContext(c).Where(q.ID.Eq(download.ID)).Updates(rollbackUpdates) resputil.HandleError(c, bizerr.Internal.K8sServiceError.Wrap(err, "submit download job failed")) return } @@ -1239,10 +1003,6 @@ func (mgr *ModelDownloadMgr) ListAllDownloads(c *gin.Context) { for i, d := range downloads { resp[i] = convertDownloadToResp(d, token) } - if err := mgr.applyDownloadUserContext(c, downloads, resp, token); err != nil { - resputil.HandleError(c, err) - return - } resputil.Success(c, resp) } @@ -1294,13 +1054,6 @@ func (mgr *ModelDownloadMgr) AdminDeleteDownload(c *gin.Context) { // a couple of times before being marked Failed. const downloadJobBackoffLimit int32 = 2 -const ( - defaultModelDownloaderImage = "ghcr.io/raids-lab/crater-model-downloader:v1.0.0" - huggingFaceHubVersion = "1.23.0" - modelScopeVersion = "1.38.1" - modelScopeHubVersion = "0.1.7" -) - func (mgr *ModelDownloadMgr) submitDownloadJob(c *gin.Context, download *model.ModelDownload, username, accessToken string) error { physicalPath := mgr.convertToPhysicalPath(download.Path) subPath := filepath.Dir(physicalPath) @@ -1332,8 +1085,7 @@ func (mgr *ModelDownloadMgr) submitDownloadJob(c *gin.Context, download *model.M }, }, Spec: corev1.PodSpec{ - RestartPolicy: corev1.RestartPolicyNever, - ImagePullSecrets: downloadImagePullSecrets(config.GetConfig().Secrets.ImagePullSecretName), + RestartPolicy: corev1.RestartPolicyNever, Containers: []corev1.Container{ { Name: "downloader", @@ -1380,17 +1132,6 @@ func (mgr *ModelDownloadMgr) submitDownloadJob(c *gin.Context, download *model.M return err } -// downloadImagePullSecrets keeps the public-image path credential-free while -// allowing deployments to use a private, internally mirrored downloader image. -// The secret itself remains an administrator-owned deployment setting; download -// requests never accept or persist registry credentials. -func downloadImagePullSecrets(secretName string) []corev1.LocalObjectReference { - if secretName == "" { - return nil - } - return []corev1.LocalObjectReference{{Name: secretName}} -} - // downloadTokenEnv injects the access token for the given source as the env var // the corresponding SDK expects, when a token is provided. func downloadTokenEnv(source model.ModelSource, accessToken string) []corev1.EnvVar { @@ -1455,27 +1196,22 @@ func (mgr *ModelDownloadMgr) getDownloadImage(_ model.ModelSource) string { if img := config.GetConfig().ModelDownload.Image; img != "" { return img } - // The public default is reproducible and can be mirrored or overridden by each deployment. - return defaultModelDownloaderImage + // 使用官方 Python 镜像作为默认值 + // 本地部署可通过配置文件指定内网镜像 + return "python:3.11-slim" } func (mgr *ModelDownloadMgr) buildDownloadCommand(download *model.ModelDownload, modelDirName string) string { - var installCmd, preflightCmd, downloadCommand, totalProbeCmd, metadataCmd string - huggingFaceEndpoint := config.GetConfig().HuggingFaceDownloadEndpoint() - disableXetCmd := "" - if shouldDisableHuggingFaceXet(download.Source, huggingFaceEndpoint) { - disableXetCmd = "export HF_HUB_DISABLE_XET=1" - } + // 清华源配置 + pypiMirror := "https://pypi.tuna.tsinghua.edu.cn/simple" + trustedHost := "--trusted-host pypi.tuna.tsinghua.edu.cn" + + var installCmd, downloadCommand, totalProbeCmd, metadataCmd string if download.Source == model.ModelSourceHuggingFace { - // The project image already contains this dependency. The fallback preserves - // compatibility with existing custom Python images used by open-source deployments. + // 1. 安装 huggingface_hub installCmd = fmt.Sprintf( - `if ! python -c 'import huggingface_hub' >/dev/null 2>&1; then - echo "huggingface_hub is missing; installing the tested fallback version" - pip install --no-cache-dir 'huggingface_hub==%s' -fi -python -c 'import huggingface_hub; print("[TOOLS] huggingface_hub=" + huggingface_hub.__version__)'`, - huggingFaceHubVersion, + "pip install -U 'huggingface_hub>=0.23.0' -i %s %s", + pypiMirror, trustedHost, ) // 2. 用 Python API snapshot_download，而不是 huggingface-cli @@ -1530,7 +1266,6 @@ PY echo "Reading repository metadata..." python - << 'PY' || true import json -import os import urllib.error import urllib.parse import urllib.request @@ -1548,10 +1283,8 @@ parameter_count = safetensors.get("total", 0) if isinstance(safetensors, dict) e owner = %q.split("/", 1)[0] avatar_url = "" for owner_type in ("organizations", "users"): - endpoint = "{}/api/{}/{}/overview".format( - os.environ.get("HF_ENDPOINT", "https://huggingface.co").rstrip("/"), - owner_type, - urllib.parse.quote(owner, safe=""), + endpoint = "https://huggingface.co/api/{}/{}/overview".format( + owner_type, urllib.parse.quote(owner, safe="") ) try: with urllib.request.urlopen(endpoint, timeout=10) as response: @@ -1581,71 +1314,13 @@ print("[META] " + json.dumps({ PY `, download.Name, repoType, download.Revision, download.Name, download.Name) } else { - // The project image already contains these dependencies. The fallback keeps - // user-supplied Python images working while avoiding an unpinned installation. + // Install the ModelScope CLI, then invoke it through an argument array so + // user-provided revisions never pass through shell parsing. installCmd = fmt.Sprintf( - `if ! python -c 'import modelscope, modelscope_hub' >/dev/null 2>&1; then - echo "ModelScope clients are missing; installing the tested fallback versions" - pip install --no-cache-dir 'modelscope==%s' 'modelscope-hub==%s' -fi -python - << 'PY' -import modelscope -import modelscope_hub -print("[TOOLS] modelscope={} modelscope_hub={}".format(modelscope.__version__, modelscope_hub.__version__)) -PY`, - modelScopeVersion, modelScopeHubVersion, + "pip install -q modelscope -i %s %s", + pypiMirror, trustedHost, ) - // Validate an explicit revision before downloading. ModelScope's legacy file - // API currently returns Files=null for an unknown branch, which otherwise - // surfaces as an opaque "NoneType is not iterable" client error. - resourcePath := "models" - if download.Category == model.DownloadCategoryDataset { - resourcePath = "datasets" - } - preflightCmd = fmt.Sprintf(` -echo "Validating ModelScope revision..." -python - << 'PY' -import json -import os -import sys -import urllib.parse -import urllib.request - -repo_id = %q -revision = %q -resource_path = %q -if revision: - endpoint = os.environ.get("MODELSCOPE_ENDPOINT", "https://modelscope.cn").rstrip("/") - url = endpoint + "/api/v1/{}/{}/revisions".format( - resource_path, urllib.parse.quote(repo_id, safe="/")) - request = urllib.request.Request(url) - token = os.environ.get("MODELSCOPE_API_TOKEN", "") - if token: - request.add_header("Authorization", "Bearer " + token) - try: - with urllib.request.urlopen(request, timeout=20) as response: - payload = json.load(response) - revision_map = (payload.get("Data") or payload.get("data") or {}).get("RevisionMap") or {} - entries = (revision_map.get("Branches") or []) + (revision_map.get("Tags") or []) - available = sorted({entry.get("Revision") for entry in entries if entry.get("Revision")}) - except Exception as error: - print("[WARN] revision validation unavailable: {}".format(error), flush=True) - else: - if available and revision not in available: - print( - "[ERROR] revision_not_found: {!r} is not available for {}; available revisions: {}".format( - revision, repo_id, ", ".join(available) - ), - file=sys.stderr, - flush=True, - ) - raise SystemExit(22) -PY -`, download.Name, download.Revision, resourcePath) - - // Invoke the CLI through an argument array so - // user-provided revisions never pass through shell parsing. var resourceFlag string if download.Category == model.DownloadCategoryDataset { resourceFlag = "--dataset" @@ -1690,6 +1365,10 @@ except Exception: PY `, download.Name, download.Revision, string(download.Category)) + resourcePath := "models" + if download.Category == model.DownloadCategoryDataset { + resourcePath = "datasets" + } metadataCmd = fmt.Sprintf(` echo "Reading repository metadata..." python - << 'PY' || true @@ -1697,7 +1376,7 @@ import json import os import urllib.request -url = os.environ.get("MODELSCOPE_ENDPOINT", "https://modelscope.cn").rstrip("/") + "/openapi/v1/%s/%s" +url = "https://modelscope.cn/openapi/v1/%s/%s" request = urllib.request.Request(url) token = os.environ.get("MODELSCOPE_API_TOKEN", "") if token: @@ -1777,9 +1456,7 @@ trap "kill $MONITOR_PID 2>/dev/null || true" EXIT return fmt.Sprintf(` set -euo pipefail -export HF_ENDPOINT=%q -%s -export MODELSCOPE_ENDPOINT=%q +export HF_ENDPOINT=https://hf-mirror.com OUT_DIR="/data/%s" export OUT_DIR mkdir -p "$OUT_DIR" @@ -1793,7 +1470,8 @@ on_error() { } trap 'on_error $LINENO' ERR -# Verify preinstalled tools, with a pinned fallback for custom legacy images. +# 安装依赖 +echo "Installing dependencies from Tsinghua mirror..." %s # 确保 Python 包路径可用 @@ -1805,8 +1483,6 @@ export PATH="/usr/local/bin:/root/.local/bin:$PATH" %s -%s - # 执行下载 START_TIME=$(date +%%s) %s @@ -1833,15 +1509,11 @@ else echo "[RESULT] size_bytes=$SIZE" fi `, - huggingFaceEndpoint, - disableXetCmd, - config.GetConfig().ModelScopeDownloadEndpoint(), modelDirName, download.Name, download.Source, installCmd, totalProbeCmd, - preflightCmd, progressScript, downloadCommand, metadataCmd, @@ -1849,11 +1521,6 @@ fi ) } -func shouldDisableHuggingFaceXet(source model.ModelSource, endpoint string) bool { - return source == model.ModelSourceHuggingFace && - strings.TrimRight(strings.TrimSpace(endpoint), "/") != "https://huggingface.co" -} - // convertToPhysicalPath 将前端路径转换为物理存储路径 func (mgr *ModelDownloadMgr) convertToPhysicalPath(frontendPath string) string { // public -> sugon-gpu-incoming @@ -1881,8 +1548,15 @@ func sanitizeModelName(name string) string { return pattern.ReplaceAllString(name, "") } -func modelDownloadStoragePath(basePath, safeName string) string { - return filepath.Join(basePath, safeName) +func modelDownloadStoragePath( + basePath, safeName string, source model.ModelSource, revision string, +) string { + revisionKey := "default" + if revision != "" { + digest := fmt.Sprintf("%x", sha256.Sum256([]byte(revision))) + revisionKey = digest[:12] + } + return filepath.Join(basePath, safeName, string(source), revisionKey) } func convertDownloadToResp(d *model.ModelDownload, token util.JWTMessage) ModelDownloadResp { @@ -1897,11 +1571,6 @@ func convertDownloadToResp(d *model.ModelDownload, token util.JWTMessage) ModelD nickname = username } - relation := ModelDownloadRelationNone - if d.CreatorID == token.UserID { - relation = ModelDownloadRelationCreator - } - return ModelDownloadResp{ ID: d.ID, Name: d.Name, @@ -1917,15 +1586,11 @@ func convertDownloadToResp(d *model.ModelDownload, token util.JWTMessage) ModelD JobName: d.JobName, CreatorID: d.CreatorID, ReferenceCount: d.ReferenceCount, - RequesterCount: d.ReferenceCount, - Requesters: []model.UserInfo{}, - Relation: relation, CreatedAt: d.CreatedAt, UpdatedAt: d.UpdatedAt, SourceUpdatedAt: d.SourceUpdatedAt, UserInfo: model.UserInfo{Username: username, Nickname: nickname}, CanManage: d.CreatorID == token.UserID || token.RolePlatform == model.RoleAdmin, - CanDelete: canDeleteDownload(token), CanViewLogs: d.CreatorID == token.UserID || token.RolePlatform == model.RoleAdmin, SourceURL: sourceURLForDownload(d), DisplayName: d.DisplayName, @@ -1943,17 +1608,15 @@ func sourceURLForDownload(download *model.ModelDownload) string { return download.SourceURL } if download.Source == model.ModelSourceHuggingFace { - endpoint := config.GetConfig().HuggingFaceDownloadEndpoint() if download.Category == model.DownloadCategoryDataset { - return endpoint + "/datasets/" + download.Name + return "https://huggingface.co/datasets/" + download.Name } - return endpoint + "/" + download.Name + return "https://huggingface.co/" + download.Name } - endpoint := config.GetConfig().ModelScopeDownloadEndpoint() if download.Category == model.DownloadCategoryDataset { - return endpoint + "/datasets/" + download.Name + return "https://modelscope.cn/datasets/" + download.Name } - return endpoint + "/models/" + download.Name + return "https://modelscope.cn/models/" + download.Name } func (mgr *ModelDownloadMgr) deleteDownloadJob(c *gin.Context, jobName string) error { @@ -1973,11 +1636,6 @@ func (mgr *ModelDownloadMgr) deleteDownloadJob(c *gin.Context, jobName string) e func (mgr *ModelDownloadMgr) deleteDownloadRecord(c *gin.Context, downloadID uint) error { db := query.Use(query.GetDB()) err := db.Transaction(func(tx *query.Query) error { - if releaseErr := service.ReleaseModelDownloadQuotaReservation( - c, tx.ModelDownload.WithContext(c).UnderlyingDB().Session(&gorm.Session{NewDB: true}), downloadID, - ); releaseErr != nil { - return releaseErr - } qUserDownload := tx.UserModelDownload if _, deleteErr := qUserDownload.WithContext(c). Where(qUserDownload.ModelDownloadID.Eq(downloadID)). diff --git a/backend/internal/handler/modeldownload_test.go b/backend/internal/handler/modeldownload_test.go index 4208b13e2..847768732 100644 --- a/backend/internal/handler/modeldownload_test.go +++ b/backend/internal/handler/modeldownload_test.go @@ -15,204 +15,13 @@ package handler import ( - "encoding/json" - "errors" - "net/http/httptest" "path/filepath" "strings" "testing" - "github.com/gin-gonic/gin" - "gorm.io/driver/sqlite" - "gorm.io/gorm" - "github.com/raids-lab/crater/dao/model" - "github.com/raids-lab/crater/dao/query" - "github.com/raids-lab/crater/internal/bizerr" - "github.com/raids-lab/crater/internal/util" ) -type testSQLStateError struct { - state string -} - -func (err testSQLStateError) Error() string { - return "database error with SQLSTATE " + err.state -} - -func (err testSQLStateError) SQLState() string { - return err.state -} - -func TestDownloadActionRevisionDistinguishesOmittedAndDefault(t *testing.T) { - var omitted DownloadActionReq - if err := json.Unmarshal([]byte(`{"token":"temporary"}`), &omitted); err != nil { - t.Fatal(err) - } - if omitted.Revision != nil { - t.Fatalf("omitted revision should preserve the failed record revision: %#v", omitted.Revision) - } - - var defaultBranch DownloadActionReq - if err := json.Unmarshal([]byte(`{"revision":""}`), &defaultBranch); err != nil { - t.Fatal(err) - } - if defaultBranch.Revision == nil || *defaultBranch.Revision != "" { - t.Fatalf("explicit empty revision should select the source default: %#v", defaultBranch.Revision) - } -} - -func TestNormalizeRetryRevision(t *testing.T) { - revision := " master " - action := DownloadActionReq{Revision: &revision} - if err := normalizeRetryRevision(&action); err != nil { - t.Fatal(err) - } - if action.Revision == nil || *action.Revision != "master" { - t.Fatalf("retry revision should be trimmed: %#v", action.Revision) - } - - tooLong := strings.Repeat("x", maxDownloadRevisionLength+1) - if err := normalizeRetryRevision(&DownloadActionReq{Revision: &tooLong}); err == nil { - t.Fatal("overlong retry revision should be rejected") - } -} - -func TestDownloadRelationshipDefaults(t *testing.T) { - creator := util.JWTMessage{UserID: 7} - download := &model.ModelDownload{CreatorID: creator.UserID} - if got := convertDownloadToResp(download, creator).Relation; got != ModelDownloadRelationCreator { - t.Fatalf("creator relation = %q", got) - } - if got := convertDownloadToResp(download, util.JWTMessage{UserID: 8}).Relation; got != ModelDownloadRelationNone { - t.Fatalf("unassociated user relation = %q", got) - } -} - -func TestModelDownloadResponseKeepsReferenceCountCompatibilityAlias(t *testing.T) { - download := &model.ModelDownload{CreatorID: 7, ReferenceCount: 3} - response := convertDownloadToResp(download, util.JWTMessage{UserID: 7}) - if response.ReferenceCount != response.RequesterCount || response.RequesterCount != 3 { - t.Fatalf("referenceCount/requesterCount = %d/%d, want 3/3", - response.ReferenceCount, response.RequesterCount) - } - - encoded, err := json.Marshal(response) - if err != nil { - t.Fatal(err) - } - for _, field := range []string{`"referenceCount":3`, `"requesterCount":3`} { - if !strings.Contains(string(encoded), field) { - t.Fatalf("response JSON %s does not contain compatibility field %s", encoded, field) - } - } -} - -func TestApplyDownloadRequestersRecordsDemandWithoutChangingPublicAccess(t *testing.T) { - responses := []ModelDownloadResp{ - {ID: 10, Relation: ModelDownloadRelationCreator}, - {ID: 11, Relation: ModelDownloadRelationNone}, - } - associations := []*model.UserModelDownload{ - {ModelDownloadID: 10, UserID: 1, User: model.User{Name: "alice", Nickname: "Alice"}}, - {ModelDownloadID: 10, UserID: 2, User: model.User{Name: "bob"}}, - {ModelDownloadID: 10, UserID: 3}, // Keep deleted users in the demand count. - {ModelDownloadID: 11, UserID: 1, User: model.User{Name: "alice", Nickname: "Alice"}}, - } - - applyDownloadRequesters(responses, associations, util.JWTMessage{UserID: 1}) - - if len(responses) != 2 { - t.Fatalf("response count = %d, want 2", len(responses)) - } - creatorResponse := responses[0] - if creatorResponse.Relation != ModelDownloadRelationCreator { - t.Fatalf("creator relation changed to %q", creatorResponse.Relation) - } - if creatorResponse.RequesterCount != 3 || len(creatorResponse.Requesters) != 2 { - t.Fatalf("requesters = %d/%d, want 3 recorded and 2 visible", creatorResponse.RequesterCount, - len(creatorResponse.Requesters)) - } - visibleRequesters := creatorResponse.Requesters - if len(visibleRequesters) != 2 { - t.Fatalf("visible requester count = %d, want 2", len(visibleRequesters)) - } - if visibleRequesters[1] != (model.UserInfo{Username: "bob", Nickname: "bob"}) { - t.Fatalf("empty nickname should fall back to username: %#v", visibleRequesters[1]) - } - requesterResponse := responses[1] - if requesterResponse.Relation != ModelDownloadRelationSubmitted || !requesterResponse.CanViewLogs { - t.Fatalf("requesting user context not applied: %#v", requesterResponse) - } -} - -func TestCheckRetryRevisionConflictIncludesSoftDeletedRecords(t *testing.T) { - db, err := gorm.Open(sqlite.Open("file:retry_revision_conflict?mode=memory&cache=shared"), &gorm.Config{ - DisableForeignKeyConstraintWhenMigrating: true, - IgnoreRelationshipsWhenMigrating: true, - }) - if err != nil { - t.Fatal(err) - } - if err := db.AutoMigrate(&model.ModelDownload{}); err != nil { - t.Fatal(err) - } - - active := model.ModelDownload{ - Name: "owner/model", Source: model.ModelSourceHuggingFace, - Category: model.DownloadCategoryModel, Revision: "old", Path: "public/Models/owner/model", - Status: model.ModelDownloadStatusFailed, CreatorID: 1, - } - softDeleted := active - softDeleted.Revision = "new" - if err := db.Create(&active).Error; err != nil { - t.Fatal(err) - } - if err := db.Create(&softDeleted).Error; err != nil { - t.Fatal(err) - } - if err := db.Delete(&softDeleted).Error; err != nil { - t.Fatal(err) - } - - ginContext, _ := gin.CreateTestContext(httptest.NewRecorder()) - revision := "new" - err = checkRetryRevisionConflict(ginContext, query.Use(db), &active, &revision) - if !errors.Is(err, bizerr.Conflict.Base) { - t.Fatalf("expected a conflict for soft-deleted revision, got %v", err) - } -} - -func TestRetryUpdateDuplicateIsConflict(t *testing.T) { - if err := retryUpdateError(gorm.ErrDuplicatedKey); !errors.Is(err, bizerr.Conflict.Base) { - t.Fatalf("duplicate retry update should be a conflict: %v", err) - } - if err := retryUpdateError(testSQLStateError{state: "23505"}); !errors.Is(err, bizerr.Conflict.Base) { - t.Fatalf("PostgreSQL unique violation should be a conflict: %v", err) - } - if err := retryUpdateError(gorm.ErrInvalidData); !errors.Is(err, bizerr.Internal.Base) { - t.Fatalf("non-duplicate retry update should remain internal: %v", err) - } -} - -func TestRestoredReadyDownloadDoesNotSubmitJob(t *testing.T) { - for _, testCase := range []struct { - status model.ModelDownloadStatus - want bool - }{ - {status: model.ModelDownloadStatusReady, want: false}, - {status: model.ModelDownloadStatusPending, want: true}, - {status: model.ModelDownloadStatusDownloading, want: true}, - {status: model.ModelDownloadStatusPaused, want: true}, - {status: model.ModelDownloadStatusFailed, want: true}, - } { - download := &model.ModelDownload{Status: testCase.status} - if got := shouldSubmitRestoredDownload(download); got != testCase.want { - t.Fatalf("status %s: shouldSubmitRestoredDownload() = %v, want %v", testCase.status, got, testCase.want) - } - } -} - func TestModelScopeDownloadCommandUsesArgumentArray(t *testing.T) { download := &model.ModelDownload{ Name: "Qwen/Qwen3-32B", @@ -227,10 +36,6 @@ func TestModelScopeDownloadCommandUsesArgumentArray(t *testing.T) { `args = ["modelscope", "download", resource_flag, repo_id]`, `args.extend(["--revision", revision])`, "subprocess.run(args, check=True)", - "revision_not_found", - "available revisions", - "modelscope==" + modelScopeVersion, - "modelscope-hub==" + modelScopeHubVersion, } { if !strings.Contains(command, expected) { t.Fatalf("download command does not contain %q", expected) @@ -239,265 +44,23 @@ func TestModelScopeDownloadCommandUsesArgumentArray(t *testing.T) { if strings.Contains(command, "modelscope download --model Qwen/Qwen3-32B --revision") { t.Fatal("download command interpolates arguments into a shell command") } - if strings.Contains(command, "pip install -q modelscope") { - t.Fatal("download command performs an unpinned runtime installation") - } - if strings.Contains(command, "%!") { - t.Fatalf("download command contains an unresolved format directive: %s", command) - } -} - -func TestShouldDisableHuggingFaceXet(t *testing.T) { - for _, testCase := range []struct { - name string - source model.ModelSource - endpoint string - want bool - }{ - { - name: "hugging face mirror", - source: model.ModelSourceHuggingFace, - endpoint: "https://hf-mirror.com", - want: true, - }, - { - name: "official hugging face endpoint", - source: model.ModelSourceHuggingFace, - endpoint: "https://huggingface.co/", - want: false, - }, - { - name: "modelscope download", - source: model.ModelSourceModelScope, - endpoint: "https://hf-mirror.com", - want: false, - }, - } { - t.Run(testCase.name, func(t *testing.T) { - if got := shouldDisableHuggingFaceXet(testCase.source, testCase.endpoint); got != testCase.want { - t.Fatalf("shouldDisableHuggingFaceXet(%q, %q) = %v, want %v", - testCase.source, testCase.endpoint, got, testCase.want) - } - }) - } } -func TestModelDownloadStoragePathUsesCanonicalShortPath(t *testing.T) { +func TestModelDownloadStoragePathSeparatesSourceAndRevision(t *testing.T) { base := "public/Models" name := "Qwen/Qwen3-32B" - path := modelDownloadStoragePath(base, name) - - if path != filepath.Join(base, name) { - t.Fatalf("unexpected canonical storage path: %s", path) - } -} - -func TestFindReadyLogicalDownloadReusesHistoricalPath(t *testing.T) { - db, err := gorm.Open(sqlite.Open("file:logical_download_reuse?mode=memory&cache=shared"), &gorm.Config{ - DisableForeignKeyConstraintWhenMigrating: true, - IgnoreRelationshipsWhenMigrating: true, - }) - if err != nil { - t.Fatal(err) - } - if err := db.AutoMigrate(&model.ModelDownload{}); err != nil { - t.Fatal(err) - } - - longPath := model.ModelDownload{ - Name: "Qwen/Qwen3-32B", Source: model.ModelSourceModelScope, - Category: model.DownloadCategoryModel, Revision: "master", - Path: "public/Models/Qwen/Qwen3-32B/modelscope/fc613b4dfd67", Status: model.ModelDownloadStatusReady, - CreatorID: 1, - } - shortPath := model.ModelDownload{ - Name: "Qwen/Qwen3-32B", Source: model.ModelSourceHuggingFace, - Category: model.DownloadCategoryModel, Revision: "main", - Path: "public/Models/Qwen/Qwen3-32B", Status: model.ModelDownloadStatusReady, - CreatorID: 2, - } - if err := db.Create(&longPath).Error; err != nil { - t.Fatal(err) - } - if err := db.Create(&shortPath).Error; err != nil { - t.Fatal(err) - } - - ginContext, _ := gin.CreateTestContext(httptest.NewRecorder()) - got, err := (&ModelDownloadMgr{}).findReadyOrOngoingDownload( - ginContext, query.Use(db), "Qwen/Qwen3-32B", model.ModelSourceModelScope, - model.DownloadCategoryModel, "master", - ) - if err != nil { - t.Fatal(err) - } - if got == nil || got.ID != longPath.ID { - t.Fatalf("findReadyOrOngoingDownload() = %#v, want exact historical record %d", got, longPath.ID) - } -} + defaultPath := modelDownloadStoragePath(base, name, model.ModelSourceModelScope, "") + huggingFacePath := modelDownloadStoragePath(base, name, model.ModelSourceHuggingFace, "") + revisionPath := modelDownloadStoragePath(base, name, model.ModelSourceModelScope, "v2") -func TestFindOngoingLogicalDownloadDoesNotReuseReadyOtherSource(t *testing.T) { - db, err := gorm.Open(sqlite.Open("file:logical_download_exact_ongoing?mode=memory&cache=shared"), &gorm.Config{ - DisableForeignKeyConstraintWhenMigrating: true, - IgnoreRelationshipsWhenMigrating: true, - }) - if err != nil { - t.Fatal(err) - } - if err := db.AutoMigrate(&model.ModelDownload{}); err != nil { - t.Fatal(err) - } - - readyOtherSource := model.ModelDownload{ - Name: "Qwen/Qwen3-32B", Source: model.ModelSourceHuggingFace, - Category: model.DownloadCategoryModel, Revision: "main", - Path: "public/Models/Qwen/Qwen3-32B", Status: model.ModelDownloadStatusReady, - CreatorID: 1, - } - ongoingExact := model.ModelDownload{ - Name: "Qwen/Qwen3-32B", Source: model.ModelSourceModelScope, - Category: model.DownloadCategoryModel, Revision: "master", - Path: "public/Models/Qwen/Qwen3-32B/modelscope/fc613b4dfd67", Status: model.ModelDownloadStatusDownloading, - CreatorID: 2, + if defaultPath != filepath.Join(base, name, "modelscope", "default") { + t.Fatalf("unexpected default storage path: %s", defaultPath) } - if err := db.Create(&readyOtherSource).Error; err != nil { - t.Fatal(err) + if defaultPath == huggingFacePath || defaultPath == revisionPath || huggingFacePath == revisionPath { + t.Fatalf("source/revision paths must be distinct: %q %q %q", defaultPath, huggingFacePath, revisionPath) } - if err := db.Create(&ongoingExact).Error; err != nil { - t.Fatal(err) - } - - ginContext, _ := gin.CreateTestContext(httptest.NewRecorder()) - got, err := (&ModelDownloadMgr{}).findReadyOrOngoingDownload( - ginContext, query.Use(db), "Qwen/Qwen3-32B", model.ModelSourceModelScope, - model.DownloadCategoryModel, "master", - ) - if err != nil { - t.Fatal(err) - } - if got == nil || got.ID != ongoingExact.ID { - t.Fatalf("findReadyOrOngoingDownload() = %#v, want exact ongoing record %d", got, ongoingExact.ID) - } - - got, err = (&ModelDownloadMgr{}).findReadyOrOngoingDownload( - ginContext, query.Use(db), "Qwen/Qwen3-32B", model.ModelSourceModelScope, - model.DownloadCategoryModel, "v2", - ) - if err != nil { - t.Fatal(err) - } - if got != nil { - t.Fatalf("cross-revision request reused download %#v", got) - } - if err := checkLogicalDownloadConflict( - ginContext, query.Use(db), "Qwen/Qwen3-32B", model.ModelSourceModelScope, - model.DownloadCategoryModel, "v2", - ); !errors.Is(err, bizerr.Conflict.Base) { - t.Fatalf("cross-revision request should conflict with canonical storage, got %v", err) - } -} - -func TestAssociateUserWithLogicalDownloadIncrementsReferenceOnce(t *testing.T) { - db, err := gorm.Open(sqlite.Open("file:logical_download_reference?mode=memory&cache=shared"), &gorm.Config{ - DisableForeignKeyConstraintWhenMigrating: true, - IgnoreRelationshipsWhenMigrating: true, - }) - if err != nil { - t.Fatal(err) - } - if err := db.AutoMigrate(&model.ModelDownload{}, &model.UserModelDownload{}); err != nil { - t.Fatal(err) - } - - download := model.ModelDownload{ - Name: "Qwen/Qwen3-32B", Source: model.ModelSourceModelScope, - Category: model.DownloadCategoryModel, Revision: "master", - Path: "public/Models/Qwen/Qwen3-32B/modelscope/fc613b4dfd67", Status: model.ModelDownloadStatusReady, - CreatorID: 1, ReferenceCount: 1, - } - if err := db.Create(&download).Error; err != nil { - t.Fatal(err) - } - if err := db.Create(&model.UserModelDownload{UserID: 1, ModelDownloadID: download.ID}).Error; err != nil { - t.Fatal(err) - } - - ginContext, _ := gin.CreateTestContext(httptest.NewRecorder()) - mgr := &ModelDownloadMgr{} - for range 2 { - if err := mgr.associateUserWithDownload(ginContext, query.Use(db), 2, download.ID); err != nil { - t.Fatal(err) - } - } - - var updated model.ModelDownload - if err := db.First(&updated, download.ID).Error; err != nil { - t.Fatal(err) - } - if updated.ReferenceCount != 2 { - t.Fatalf("reference count = %d, want 2 unique users", updated.ReferenceCount) - } - var associations int64 - if err := db.Model(&model.UserModelDownload{}). - Where("model_download_id = ? AND user_id = ?", download.ID, 2). - Count(&associations).Error; err != nil { - t.Fatal(err) - } - if associations != 1 { - t.Fatalf("user association count = %d, want 1", associations) - } -} - -func TestLogicalDownloadConflictIncludesOtherSourceAndSoftDeletedRows(t *testing.T) { - db, err := gorm.Open(sqlite.Open("file:logical_download_conflict?mode=memory&cache=shared"), &gorm.Config{ - DisableForeignKeyConstraintWhenMigrating: true, - IgnoreRelationshipsWhenMigrating: true, - }) - if err != nil { - t.Fatal(err) - } - if err := db.AutoMigrate(&model.ModelDownload{}); err != nil { - t.Fatal(err) - } - - download := model.ModelDownload{ - Name: "Qwen/Qwen3-32B", Source: model.ModelSourceModelScope, - Category: model.DownloadCategoryModel, Revision: "master", - Path: "public/Models/Qwen/Qwen3-32B/modelscope/fc613b4dfd67", Status: model.ModelDownloadStatusFailed, - CreatorID: 1, - } - if err := db.Create(&download).Error; err != nil { - t.Fatal(err) - } - if err := db.Delete(&download).Error; err != nil { - t.Fatal(err) - } - - ginContext, _ := gin.CreateTestContext(httptest.NewRecorder()) - err = checkLogicalDownloadConflict( - ginContext, query.Use(db), "Qwen/Qwen3-32B", model.ModelSourceHuggingFace, - model.DownloadCategoryModel, "main", - ) - if !errors.Is(err, bizerr.Conflict.Base) { - t.Fatalf("expected cross-source soft-deleted conflict, got %v", err) - } - - if err := checkLogicalDownloadConflict( - ginContext, query.Use(db), "Qwen/Qwen3-32B", model.ModelSourceModelScope, - model.DownloadCategoryModel, "master", - ); err != nil { - t.Fatalf("exact historical identity should be retried in place, got %v", err) - } -} - -func TestDownloadImagePullSecrets(t *testing.T) { - if got := downloadImagePullSecrets(""); got != nil { - t.Fatalf("public downloader image should not require pull secrets: %#v", got) - } - - got := downloadImagePullSecrets("internal-registry") - if len(got) != 1 || got[0].Name != "internal-registry" { - t.Fatalf("private downloader image should use the configured pull secret: %#v", got) + if revisionPath != modelDownloadStoragePath(base, name, model.ModelSourceModelScope, "v2") { + t.Fatal("revision storage path must be deterministic") } } @@ -525,18 +88,3 @@ func TestTruncateDownloadLogTail(t *testing.T) { t.Fatalf("unexpected truncated log tail: %q", truncated) } } - -func TestDownloadRecordDeletionIsAdminOnly(t *testing.T) { - download := &model.ModelDownload{CreatorID: 42} - creator := util.JWTMessage{UserID: 42, RolePlatform: model.RoleUser} - creatorResponse := convertDownloadToResp(download, creator) - if !creatorResponse.CanManage || creatorResponse.CanDelete || canDeleteDownload(creator) { - t.Fatalf("creator permissions = %#v", creatorResponse) - } - - admin := util.JWTMessage{UserID: 7, RolePlatform: model.RoleAdmin} - adminResponse := convertDownloadToResp(download, admin) - if !adminResponse.CanManage || !adminResponse.CanDelete || !canDeleteDownload(admin) { - t.Fatalf("admin permissions = %#v", adminResponse) - } -} diff --git a/backend/internal/handler/node.go b/backend/internal/handler/node.go index 70adc796b..26ec510ba 100644 --- a/backend/internal/handler/node.go +++ b/backend/internal/handler/node.go @@ -86,6 +86,7 @@ func (mgr *NodeMgr) RegisterProtected(g *gin.RouterGroup) { g.GET("/:name/gpu", mgr.ListNodeGPUInfo) } +//nolint:dupl // Router registration blocks naturally look similar across managers. func (mgr *NodeMgr) RegisterAdmin(g *gin.RouterGroup) { g.GET("", mgr.ListNode) g.GET("/:name/pods", mgr.AdminGetPodsForNode) diff --git a/backend/internal/handler/system_config.go b/backend/internal/handler/system_config.go index e1564faea..109387049 100644 --- a/backend/internal/handler/system_config.go +++ b/backend/internal/handler/system_config.go @@ -2,8 +2,6 @@ package handler import ( "fmt" - "slices" - "strings" "time" "github.com/gin-gonic/gin" @@ -11,11 +9,12 @@ import ( "gorm.io/gorm/clause" "k8s.io/klog/v2" + "strings" + "github.com/raids-lab/crater/dao/query" "github.com/raids-lab/crater/internal/bizerr" "github.com/raids-lab/crater/internal/resputil" "github.com/raids-lab/crater/internal/service" - "github.com/raids-lab/crater/internal/util" "github.com/raids-lab/crater/pkg/cronjob" "github.com/raids-lab/crater/pkg/prequeuewatcher" ) @@ -47,7 +46,6 @@ func (mgr *SystemConfigMgr) GetName() string { return mgr.name func (mgr *SystemConfigMgr) RegisterPublic(_ *gin.RouterGroup) {} func (mgr *SystemConfigMgr) RegisterProtected(g *gin.RouterGroup) { g.GET("/billing", mgr.GetBillingStatus) - g.GET("/model-download-limit", mgr.GetModelDownloadLimitConfig) } func (mgr *SystemConfigMgr) RegisterAdmin(g *gin.RouterGroup) { @@ -61,8 +59,6 @@ func (mgr *SystemConfigMgr) RegisterAdmin(g *gin.RouterGroup) { g.PUT("/gpu-analysis", mgr.SetGpuAnalysisStatus) g.GET("/prequeue", mgr.GetPrequeueConfig) g.PUT("/prequeue", mgr.UpdatePrequeueConfig) - g.GET("/model-download-limit", mgr.GetAdminModelDownloadLimitConfig) - g.PUT("/model-download-limit", mgr.UpdateModelDownloadLimitConfig) g.GET("/billing", mgr.GetBillingStatus) g.PUT("/billing", mgr.SetBillingStatus) @@ -112,30 +108,6 @@ type UpdatePrequeueConfigReq struct { PrequeueCandidateSize *int64 `json:"prequeueCandidateSize" binding:"required,gt=0"` } -type ModelDownloadLimitConfigResp struct { - Enabled bool `json:"enabled"` - MaxConcurrent int64 `json:"maxConcurrent"` - WindowHours int64 `json:"windowHours"` - MaxSuccessfulDownloads int64 `json:"maxSuccessfulDownloads"` - Exempt bool `json:"exempt"` -} - -type AdminModelDownloadLimitConfigResp struct { - Enabled bool `json:"enabled"` - MaxConcurrent int64 `json:"maxConcurrent"` - WindowHours int64 `json:"windowHours"` - MaxSuccessfulDownloads int64 `json:"maxSuccessfulDownloads"` - WhitelistUserIDs []uint `json:"whitelistUserIds"` -} - -type UpdateModelDownloadLimitConfigReq struct { - Enabled *bool `json:"enabled" binding:"required"` - MaxConcurrent *int64 `json:"maxConcurrent" binding:"required,gt=0"` - WindowHours *int64 `json:"windowHours" binding:"required,gt=0"` - MaxSuccessfulDownloads *int64 `json:"maxSuccessfulDownloads" binding:"required,gt=0"` - WhitelistUserIDs []uint `json:"whitelistUserIds"` -} - type BillingStatusResp struct { FeatureEnabled bool `json:"featureEnabled"` Active bool `json:"active"` @@ -382,78 +354,6 @@ func (mgr *SystemConfigMgr) UpdatePrequeueConfig(c *gin.Context) { resputil.Success(c, "Prequeue configuration updated successfully") } -// GetModelDownloadLimitConfig godoc -// -// @Summary 获取模型与数据集下载额度 -// @Description 获取当前用户的并发任务上限、滚动窗口成功下载上限和白名单豁免状态 -// @Tags SystemConfig -// @Produce json -// @Security Bearer -// @Success 200 {object} resputil.Response[ModelDownloadLimitConfigResp] -// @Router /v1/system-config/model-download-limit [get] -func (mgr *SystemConfigMgr) GetModelDownloadLimitConfig(c *gin.Context) { - cfg, err := mgr.service.GetModelDownloadLimitConfig(c.Request.Context()) - if err != nil { - resputil.HandleError(c, bizerr.Internal.DatabaseError.Wrap(err, "get model download limit config failed")) - return - } - resputil.Success(c, ModelDownloadLimitConfigResp{ - Enabled: cfg.Enabled, MaxConcurrent: cfg.MaxConcurrent, - WindowHours: cfg.WindowHours, MaxSuccessfulDownloads: cfg.MaxSuccessfulDownloads, - Exempt: slices.Contains(cfg.WhitelistUserIDs, util.GetToken(c).UserID), - }) -} - -// GetAdminModelDownloadLimitConfig godoc -// -// @Summary 管理员获取模型与数据集下载额度 -// @Description 获取下载额度和白名单用户 ID -// @Tags SystemConfig -// @Produce json -// @Security Bearer -// @Success 200 {object} resputil.Response[AdminModelDownloadLimitConfigResp] -// @Router /v1/admin/system-config/model-download-limit [get] -func (mgr *SystemConfigMgr) GetAdminModelDownloadLimitConfig(c *gin.Context) { - cfg, err := mgr.service.GetModelDownloadLimitConfig(c.Request.Context()) - if err != nil { - resputil.HandleError(c, bizerr.Internal.DatabaseError.Wrap(err, "get model download limit config failed")) - return - } - resputil.Success(c, AdminModelDownloadLimitConfigResp{ - Enabled: cfg.Enabled, MaxConcurrent: cfg.MaxConcurrent, - WindowHours: cfg.WindowHours, MaxSuccessfulDownloads: cfg.MaxSuccessfulDownloads, - WhitelistUserIDs: cfg.WhitelistUserIDs, - }) -} - -// UpdateModelDownloadLimitConfig godoc -// -// @Summary 更新模型与数据集下载额度 -// @Description 配置所有用户的并发任务上限、滚动窗口成功下载上限和豁免白名单 -// @Tags SystemConfig -// @Accept json -// @Produce json -// @Security Bearer -// @Param data body UpdateModelDownloadLimitConfigReq true "下载额度配置" -// @Success 200 {object} resputil.Response[string] -// @Router /v1/admin/system-config/model-download-limit [put] -func (mgr *SystemConfigMgr) UpdateModelDownloadLimitConfig(c *gin.Context) { - var req UpdateModelDownloadLimitConfigReq - if err := c.ShouldBindJSON(&req); err != nil { - resputil.HandleError(c, bizerr.BadRequest.ParameterError.Wrap(err, "invalid model download limit config")) - return - } - if err := mgr.service.UpdateModelDownloadLimitConfig(c.Request.Context(), service.ModelDownloadLimitConfig{ - Enabled: *req.Enabled, MaxConcurrent: *req.MaxConcurrent, - WindowHours: *req.WindowHours, MaxSuccessfulDownloads: *req.MaxSuccessfulDownloads, - WhitelistUserIDs: req.WhitelistUserIDs, - }); err != nil { - resputil.HandleError(c, bizerr.Internal.DatabaseError.Wrap(err, "update model download limit config failed")) - return - } - resputil.Success(c, "Model download limit configuration updated") -} - func (mgr *SystemConfigMgr) GetBillingStatus(c *gin.Context) { if mgr.billingService == nil { resputil.Error(c, "billing service is not initialized", resputil.ServiceError) diff --git a/backend/internal/handler/system_config_test.go b/backend/internal/handler/system_config_test.go deleted file mode 100644 index 0add07fdd..000000000 --- a/backend/internal/handler/system_config_test.go +++ /dev/null @@ -1,101 +0,0 @@ -// Copyright 2026 The Crater Project Team, RAIDS-Lab -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package handler - -import ( - "encoding/json" - "net/http" - "net/http/httptest" - "slices" - "testing" - - "github.com/gin-gonic/gin" - "gorm.io/driver/sqlite" - "gorm.io/gorm" - - "github.com/raids-lab/crater/dao/model" - "github.com/raids-lab/crater/dao/query" - "github.com/raids-lab/crater/internal/service" - "github.com/raids-lab/crater/internal/util" -) - -func TestModelDownloadLimitConfigRoutesHideWhitelistFromProtectedUsers(t *testing.T) { - db, err := gorm.Open(sqlite.Open("file:model_download_limit_handlers?mode=memory&cache=shared"), &gorm.Config{}) - if err != nil { - t.Fatal(err) - } - if err := db.AutoMigrate(&model.SystemConfig{}, &model.PrequeueConfig{}); err != nil { - t.Fatal(err) - } - configService := service.NewConfigService(query.Use(db)) - if err := configService.UpdateModelDownloadLimitConfig(t.Context(), service.ModelDownloadLimitConfig{ - Enabled: true, MaxConcurrent: 1, WindowHours: 2, MaxSuccessfulDownloads: 5, - WhitelistUserIDs: []uint{7, 9}, - }); err != nil { - t.Fatal(err) - } - - mgr := &SystemConfigMgr{service: configService} - router := gin.New() - router.Use(func(c *gin.Context) { - util.SetJWTContext(c, util.JWTMessage{UserID: 7, RolePlatform: model.RoleAdmin}) - c.Next() - }) - mgr.RegisterProtected(router.Group("/v1/system-config")) - mgr.RegisterAdmin(router.Group("/v1/admin/system-config")) - - protectedData := requestModelDownloadLimitConfig(t, router, "/v1/system-config/model-download-limit") - if _, exposed := protectedData["whitelistUserIds"]; exposed { - t.Fatal("protected model download config must not expose the full whitelist") - } - var exempt bool - if err := json.Unmarshal(protectedData["exempt"], &exempt); err != nil { - t.Fatal(err) - } - if !exempt { - t.Fatal("whitelisted current user should be marked exempt") - } - - adminData := requestModelDownloadLimitConfig(t, router, "/v1/admin/system-config/model-download-limit") - if _, exposed := adminData["exempt"]; exposed { - t.Fatal("admin model download config should return the whitelist instead of a current-user exemption") - } - var whitelistUserIDs []uint - if err := json.Unmarshal(adminData["whitelistUserIds"], &whitelistUserIDs); err != nil { - t.Fatal(err) - } - if !slices.Equal(whitelistUserIDs, []uint{7, 9}) { - t.Fatalf("admin whitelist = %v, want [7 9]", whitelistUserIDs) - } -} - -func requestModelDownloadLimitConfig( - t *testing.T, router http.Handler, path string, -) map[string]json.RawMessage { - t.Helper() - recorder := httptest.NewRecorder() - request := httptest.NewRequest(http.MethodGet, path, http.NoBody) - router.ServeHTTP(recorder, request) - if recorder.Code != http.StatusOK { - t.Fatalf("GET %s returned HTTP %d: %s", path, recorder.Code, recorder.Body.String()) - } - var response struct { - Data map[string]json.RawMessage `json:"data"` - } - if err := json.Unmarshal(recorder.Body.Bytes(), &response); err != nil { - t.Fatal(err) - } - return response.Data -} diff --git a/backend/internal/handler/vcjob/agent_insights.go b/backend/internal/handler/vcjob/agent_insights.go new file mode 100644 index 000000000..d5fba5458 --- /dev/null +++ b/backend/internal/handler/vcjob/agent_insights.go @@ -0,0 +1,345 @@ +package vcjob + +import ( + "context" + "errors" + "fmt" + "regexp" + "strings" + + "gorm.io/gorm" + v1 "k8s.io/api/core/v1" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "sigs.k8s.io/controller-runtime/pkg/client" + "volcano.sh/apis/pkg/apis/batch/v1alpha1" + + "github.com/raids-lab/crater/dao/model" + "github.com/raids-lab/crater/dao/query" + "github.com/raids-lab/crater/internal/bizerr" + "github.com/raids-lab/crater/internal/handler" + "github.com/raids-lab/crater/internal/util" + "github.com/raids-lab/crater/pkg/config" + "github.com/raids-lab/crater/pkg/crclient" + "github.com/raids-lab/crater/pkg/monitor" +) + +//nolint:gochecknoinits // The agent package discovers this facade through handler registration. +func init() { + handler.RegisterJobInsightReaderFactory(NewAgentJobInsightReader) +} + +type agentJobInsightReader struct { + mgr *VolcanojobMgr +} + +func NewAgentJobInsightReader(conf *handler.RegisterConfig) handler.JobInsightReader { + return &agentJobInsightReader{mgr: NewVolcanojobMgr(conf).(*VolcanojobMgr)} +} + +func (r *agentJobInsightReader) FindScopedJob( + ctx context.Context, + token util.JWTMessage, + jobName string, +) (*model.Job, error) { + job, err := getJob(ctx, jobName, &token) + if err == nil { + return job, nil + } + + if isRecordNotFound(err) { + k := query.Kaniko + buildQuery := k.WithContext(ctx).Where(k.ImagePackName.Eq(jobName)) + if token.RolePlatform != model.RoleAdmin { + buildQuery = buildQuery.Where(k.UserID.Eq(token.UserID)) + } + if _, buildErr := buildQuery.First(); buildErr == nil { + return nil, bizerr.BadRequest.ParameterError.New( + fmt.Sprintf( + "%q is an image build, not a platform job; use get_image_build_detail first, "+ + "then inspect its pod with k8s_get_pod_logs / k8s_get_events", + jobName, + ), + ) + } + } + return nil, bizerr.NotFound.DataBaseNotFound.Wrap(err, "job not found") +} + +func (r *agentJobInsightReader) BuildJobDetail(job *model.Job) any { + var profileData *monitor.ProfileData + if job.ProfileData != nil { + profileData = job.ProfileData.Data() + } + + var scheduleData *model.ScheduleData + if job.ScheduleData != nil { + scheduleData = job.ScheduleData.Data() + } + + var terminatedStates []v1.ContainerStateTerminated + if job.TerminatedStates != nil { + terminatedStates = job.TerminatedStates.Data() + } + + scheduleType := model.ScheduleTypeNormal + if job.ScheduleType != nil { + scheduleType = *job.ScheduleType + } + + namespace := "" + if vcjob := job.Attributes.Data(); vcjob != nil { + namespace = vcjob.Namespace + } + + return JobDetailResp{ + Name: job.Name, + Namespace: namespace, + Username: job.User.Name, + Nickname: job.User.Nickname, + UserInfo: model.UserInfo{ + Username: job.User.Name, + Nickname: job.User.Nickname, + }, + JobName: job.JobName, + JobType: job.JobType, + ScheduleType: scheduleType, + WaitingToleranceSeconds: job.WaitingToleranceSeconds, + Queue: job.Account.Nickname, + Status: job.Status, + Resources: job.Resources.Data(), + ProfileData: profileData, + ScheduleData: scheduleData, + Events: getStoredJobEvents(job), + TerminatedStates: terminatedStates, + CreationTimestamp: metav1.NewTime(job.CreationTimestamp), + RunningTimestamp: metav1.NewTime(job.RunningTimestamp), + CompletedTimestamp: metav1.NewTime(job.CompletedTimestamp), + } +} + +func (r *agentJobInsightReader) GetJobEvents( + ctx context.Context, + token util.JWTMessage, + jobName string, +) (any, error) { + job, err := r.FindScopedJob(ctx, token, jobName) + if err != nil { + return nil, err + } + vcjob := job.Attributes.Data() + storedEvents := getStoredJobEvents(job) + if vcjob == nil { + if len(storedEvents) > 0 { + return storedEvents, nil + } + return nil, bizerr.NotFound.DataBaseNotFound.New("job attributes not found") + } + + events, err := r.listLiveJobEvents(ctx, vcjob) + if err != nil { + if len(storedEvents) > 0 { + return storedEvents, nil + } + return nil, err + } + if len(events) == 0 && len(storedEvents) > 0 { + return storedEvents, nil + } + return events, nil +} + +func (r *agentJobInsightReader) GetJobLog( + ctx context.Context, + token util.JWTMessage, + jobName string, + tailLines int64, + keyword string, +) (map[string]string, error) { + job, err := r.FindScopedJob(ctx, token, jobName) + if err != nil { + return nil, err + } + return r.readJobLogPayload(ctx, job, tailLines, keyword) +} + +func (r *agentJobInsightReader) GetDiagnosticContext( + ctx context.Context, + token util.JWTMessage, + jobName string, + includeLog bool, + tailLines int64, +) (handler.JobContextResp, error) { + job, err := r.FindScopedJob(ctx, token, jobName) + if err != nil { + return handler.JobContextResp{}, err + } + + resp := handler.JobContextResp{} + resp.Meta.Name = job.Name + resp.Meta.JobName = job.JobName + resp.Meta.Namespace = getAgentJobNamespace(job) + resp.Meta.User = job.User.Name + resp.Meta.Queue = job.Account.Nickname + resp.Meta.JobType = job.JobType + resp.Meta.Status = job.Status + resp.Meta.CreationTimestamp = job.CreationTimestamp + resp.Meta.RunningTimestamp = job.RunningTimestamp + resp.Meta.CompletedTimestamp = job.CompletedTimestamp + if job.Nodes.Data() != nil { + resp.Meta.Nodes = job.Nodes.Data() + } + resp.Meta.Resources = job.Resources.Data() + + if job.ProfileData != nil { + resp.DB.ProfileData = job.ProfileData.Data() + } + if job.ScheduleData != nil { + resp.DB.ScheduleData = job.ScheduleData.Data() + } + if events, eventsErr := r.GetJobEvents(ctx, token, jobName); eventsErr == nil { + if typedEvents, ok := events.([]v1.Event); ok { + resp.DB.Events = typedEvents + } + } else if job.Events != nil { + resp.DB.Events = job.Events.Data() + } + if job.TerminatedStates != nil { + resp.DB.TerminatedStates = job.TerminatedStates.Data() + } + + if includeLog { + logPayload, logErr := r.readJobLogPayload(ctx, job, tailLines, "") + if logErr != nil { + return handler.JobContextResp{}, logErr + } + resp.Log.Container = logPayload["container"] + resp.Log.Tail = logPayload["log"] + } + return resp, nil +} + +func (r *agentJobInsightReader) listLiveJobEvents(ctx context.Context, vcjob *v1alpha1.Job) ([]v1.Event, error) { + jobEvents, err := r.mgr.kubeClient.CoreV1().Events(vcjob.Namespace).List(ctx, metav1.ListOptions{ + FieldSelector: fmt.Sprintf("involvedObject.name=%s", vcjob.Name), + TypeMeta: metav1.TypeMeta{Kind: "Job", APIVersion: "batch.volcano.sh/v1alpha1"}, + }) + if err != nil { + return nil, bizerr.Internal.K8sServiceError.Wrap(err, "failed to list job events") + } + events := jobEvents.Items + + baseURL, ok := vcjob.Labels[crclient.LabelKeyBaseURL] + if !ok || baseURL == "" { + return events, nil + } + + podList := &v1.PodList{} + labels := client.MatchingLabels{crclient.LabelKeyBaseURL: baseURL} + if err := r.mgr.client.List(ctx, podList, client.InNamespace(vcjob.Namespace), labels); err != nil { + return nil, bizerr.Internal.K8sServiceError.Wrap(err, "failed to list job pods") + } + + containsPodEvents := false + for i := range podList.Items { + pod := &podList.Items[i] + podEvents, err := r.mgr.kubeClient.CoreV1().Events(vcjob.Namespace).List(ctx, metav1.ListOptions{ + FieldSelector: fmt.Sprintf("involvedObject.name=%s", pod.Name), + TypeMeta: metav1.TypeMeta{Kind: "Pod"}, + }) + if err != nil { + return nil, bizerr.Internal.K8sServiceError.Wrap(err, "failed to list pod events") + } + if len(podEvents.Items) > 0 && !containsPodEvents { + containsPodEvents = true + events = []v1.Event{} + } + events = append(events, podEvents.Items...) + } + return events, nil +} + +func (r *agentJobInsightReader) readJobLogPayload( + ctx context.Context, + job *model.Job, + tailLines int64, + keyword string, +) (map[string]string, error) { + if tailLines <= 0 { + tailLines = 100 + } + + namespace := getAgentJobNamespace(job) + labelSelector := fmt.Sprintf("%s=%s", crclient.LabelKeyBaseURL, job.JobName) + if vcjob := job.Attributes.Data(); vcjob != nil { + if labelVal, ok := vcjob.Labels[crclient.LabelKeyBaseURL]; ok && labelVal != "" { + labelSelector = fmt.Sprintf("%s=%s", crclient.LabelKeyBaseURL, labelVal) + } + } + + podList, podErr := r.mgr.kubeClient.CoreV1().Pods(namespace).List(ctx, metav1.ListOptions{ + LabelSelector: labelSelector, + }) + if podErr != nil || len(podList.Items) == 0 { + return map[string]string{"log": "Pod not found or no live logs available."}, nil + } + + pod := podList.Items[0] + containerName := "" + if len(pod.Spec.Containers) > 0 { + containerName = pod.Spec.Containers[0].Name + } + + logBytes, logErr := r.mgr.kubeClient.CoreV1().Pods(namespace).GetLogs(pod.Name, &v1.PodLogOptions{ + Container: containerName, + TailLines: &tailLines, + }).DoRaw(ctx) + if logErr != nil { + return map[string]string{"log": fmt.Sprintf("Failed to retrieve logs: %v", logErr)}, nil + } + + logContent, filterErr := filterAgentJobLogByKeyword(string(logBytes), keyword) + if filterErr != nil { + return nil, filterErr + } + + payload := map[string]string{ + "podName": pod.Name, + "container": containerName, + "log": logContent, + } + if keyword != "" { + payload["keyword"] = keyword + } + return payload, nil +} + +func getAgentJobNamespace(job *model.Job) string { + if job != nil { + if vcjob := job.Attributes.Data(); vcjob != nil && vcjob.Namespace != "" { + return vcjob.Namespace + } + } + return config.GetConfig().Namespaces.Job +} + +func filterAgentJobLogByKeyword(logContent, keyword string) (string, error) { + if keyword == "" { + return logContent, nil + } + re, err := regexp.Compile(keyword) + if err != nil { + return "", bizerr.BadRequest.ParameterError.Wrap(err, "invalid keyword regex") + } + lines := strings.Split(logContent, "\n") + matched := make([]string, 0, len(lines)) + for _, line := range lines { + if re.MatchString(line) { + matched = append(matched, line) + } + } + return strings.Join(matched, "\n"), nil +} + +func isRecordNotFound(err error) bool { + return errors.Is(err, gorm.ErrRecordNotFound) || strings.Contains(strings.ToLower(err.Error()), "record not found") +} diff --git a/backend/internal/handler/vcjob/agent_submit.go b/backend/internal/handler/vcjob/agent_submit.go new file mode 100644 index 000000000..b594938c0 --- /dev/null +++ b/backend/internal/handler/vcjob/agent_submit.go @@ -0,0 +1,1060 @@ +package vcjob + +import ( + "context" + "encoding/json" + "fmt" + "log" + "strings" + + "github.com/google/uuid" + v1 "k8s.io/api/core/v1" + "k8s.io/apimachinery/pkg/api/resource" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/util/intstr" + "k8s.io/utils/ptr" + batch "volcano.sh/apis/pkg/apis/batch/v1alpha1" + bus "volcano.sh/apis/pkg/apis/bus/v1alpha1" + + "github.com/raids-lab/crater/dao/model" + "github.com/raids-lab/crater/internal/bizerr" + "github.com/raids-lab/crater/internal/handler" + "github.com/raids-lab/crater/internal/util" + "github.com/raids-lab/crater/pkg/aitaskctl" + "github.com/raids-lab/crater/pkg/config" + "github.com/raids-lab/crater/pkg/crclient" + "github.com/raids-lab/crater/pkg/utils" + "github.com/raids-lab/crater/pkg/vcqueue" +) + +//nolint:gochecknoinits // Register agent job submitter factory alongside the handler. +func init() { + handler.RegisterJobMutationSubmitterFactory(NewAgentJobSubmitter) +} + +// agentJobSubmitter implements handler.JobMutationSubmitter for the agent +// service. It reuses the same vcjob helper functions but returns errors +// directly instead of writing HTTP responses via resputil. +type agentJobSubmitter struct { + mgr *VolcanojobMgr +} + +func agentSubmitErrorf(format string, args ...any) error { + return bizerr.Internal.ServiceError.New(fmt.Sprintf(strings.ReplaceAll(format, "%w", "%v"), args...)) +} + +func buildAgentParallelTaskSpecs( + req *CreateTensorflowReq, + baseAffinity *v1.Affinity, + baseTolerations []v1.Toleration, + volumes []v1.Volume, + volumeMounts []v1.VolumeMount, + envs []v1.EnvVar, + labels map[string]string, + podAnnotations map[string]string, + jobType CraterJobType, +) (tasks []batch.TaskSpec, minAvailable int32) { + tasks = make([]batch.TaskSpec, len(req.Tasks)) + minAvailable = 0 + for idx := range req.Tasks { + task := &req.Tasks[idx] + taskAffinity := GenerateArchitectureNodeAffinity(task.Image, baseAffinity) + ports := make([]v1.ContainerPort, len(task.Ports)) + for portIdx, port := range task.Ports { + ports[portIdx] = v1.ContainerPort{ + ContainerPort: port.Port, + Name: port.Name, + Protocol: v1.ProtocolTCP, + } + } + + podSpec := generatePodSpecForParallelJob( + task, + taskAffinity, + baseTolerations, + volumes, + volumeMounts, + envs, + ports, + req.CpuPinningEnabled, + ) + + taskSpec := batch.TaskSpec{ + Name: task.Name, + Replicas: task.Replicas, + Template: v1.PodTemplateSpec{ + ObjectMeta: metav1.ObjectMeta{ + Labels: labels, + Annotations: podAnnotations, + }, + Spec: podSpec, + }, + } + + applyAgentParallelTaskPolicy(&taskSpec, jobType) + minAvailable += task.Replicas + tasks[idx] = taskSpec + } + return tasks, minAvailable +} + +func applyAgentParallelTaskPolicy(taskSpec *batch.TaskSpec, jobType CraterJobType) { + switch jobType { + case CraterJobTypeTensorflow: + if taskSpec.Name == volcanoTaskWorker { + taskSpec.Policies = []batch.LifecyclePolicy{ + { + Action: bus.CompleteJobAction, + Event: bus.TaskCompletedEvent, + }, + } + } + case CraterJobTypePytorch: + switch taskSpec.Name { + case volcanoTaskMaster: + taskSpec.Policies = []batch.LifecyclePolicy{ + { + Action: bus.CompleteJobAction, + Event: bus.TaskCompletedEvent, + }, + { + Action: bus.TerminateJobAction, + Event: bus.PodFailedEvent, + }, + } + case volcanoTaskWorker: + taskSpec.Template.Spec.RestartPolicy = v1.RestartPolicyOnFailure + } + } +} + +// NewAgentJobSubmitter creates a JobMutationSubmitter backed by VolcanojobMgr +// internals. We deliberately wire the full service set (configService, +// queueQuotaSvc, prequeueWatcher, billingService) so that scheduling helpers +// such as resolveJobScheduleMetadata work the same way they do on the normal +// /v1/vcjobs/* HTTP path; agent-submitted jobs should not silently bypass +// prequeue / backfill scheduling metadata. +func NewAgentJobSubmitter(conf *handler.RegisterConfig) handler.JobMutationSubmitter { + return &agentJobSubmitter{ + mgr: &VolcanojobMgr{ + name: "vcjobs", + client: conf.Client, + config: conf.KubeConfig, + kubeClient: conf.KubeClient, + imagePacker: conf.ImagePacker, + imageRegistry: conf.ImageRegistry, + serviceManager: conf.ServiceManager, + configService: conf.ConfigService, + queueQuotaSvc: conf.PrequeueService, + prequeueWatcher: conf.PrequeueWatcher, + billingService: conf.BillingService, + }, + } +} + +func (s *agentJobSubmitter) preCheckJobCreate( + ctx context.Context, + token util.JWTMessage, + scheduleType model.ScheduleType, + requireInteractiveLimit bool, +) error { + if s.mgr.billingService != nil { + if err := s.mgr.billingService.OnJobCreateCheck(ctx, token.UserID, token.AccountID, &scheduleType); err != nil { + return err + } + } + if requireInteractiveLimit { + if err := aitaskctl.CheckInteractiveLimitBeforeCreate(ctx, token.UserID, token.AccountID); err != nil { + return agentSubmitErrorf("interactive job limit reached: %v", err) + } + } + exceededResources, err := aitaskctl.CheckResourcesBeforeCreateJob(ctx, token.UserID, token.AccountID) + if err != nil { + return agentSubmitErrorf("failed to check resources: %w", err) + } + if len(exceededResources) > 0 { + return agentSubmitErrorf("resource quota exceeded: %v", exceededResources) + } + return nil +} + +func (s *agentJobSubmitter) prepareJobCreate( + ctx context.Context, + token util.JWTMessage, + scheduleType model.ScheduleType, + requireInteractiveLimit bool, + alertEnabled bool, +) error { + if err := s.preCheckJobCreate(ctx, token, scheduleType, requireInteractiveLimit); err != nil { + return err + } + if err := vcqueue.EnsureAccountQueueExists(ctx, s.mgr.client, token, token.AccountID); err != nil { + return agentSubmitErrorf("failed to ensure account queue exists: %v", err) + } + if err := vcqueue.EnsureUserQueueExists(ctx, s.mgr.client, token, token.AccountID, token.UserID); err != nil { + return agentSubmitErrorf("failed to ensure user queue exists: %v", err) + } + if alertEnabled && !utils.CheckUserEmail(ctx, token.UserID) { + return agentSubmitErrorf("email not verified") + } + return nil +} + +func (s *agentJobSubmitter) DeleteJob(ctx context.Context, token util.JWTMessage, jobName string) (any, error) { + jobName = strings.TrimSpace(jobName) + if jobName == "" { + return nil, agentSubmitErrorf("job_name is required") + } + + jobRecord, err := getJob(ctx, jobName, &token) + if err != nil { + return nil, agentSubmitErrorf("job not found or access denied: %w", err) + } + plan, err := s.mgr.buildDeleteJobPlan(ctx, jobRecord) + if err != nil { + return nil, err + } + if err := s.mgr.applyDeleteJobPlan(ctx, jobRecord, plan); err != nil { + return nil, err + } + if err := s.mgr.deleteClusterJob(ctx, plan); err != nil { + return nil, err + } + s.mgr.notifyDeletedPrequeue(plan.shouldDeleteRecord) + + return map[string]any{ + "jobName": jobRecord.JobName, + "status": "deleted", + "deletedRecord": plan.shouldDeleteRecord, + "deletedClusterJob": plan.shouldDeleteJob, + }, nil +} + +func (s *agentJobSubmitter) StopJob(ctx context.Context, token util.JWTMessage, jobName string) (any, error) { + result, err := s.DeleteJob(ctx, token, jobName) + if err != nil { + return nil, err + } + if resultMap, ok := result.(map[string]any); ok { + resultMap["status"] = "stopped" + } + return result, nil +} + +//nolint:gocyclo // Resubmit sanitizes cloned Volcano jobs and applies optional resource overrides. +func (s *agentJobSubmitter) ResubmitJob(ctx context.Context, token util.JWTMessage, rawReq json.RawMessage) (any, error) { + var req struct { + JobName string `json:"job_name"` + Name *string `json:"name"` + CPU *string `json:"cpu"` + Memory *string `json:"memory"` + GPUCount *int `json:"gpu_count"` + GPUModel *string `json:"gpu_model"` + } + if err := json.Unmarshal(rawReq, &req); err != nil { + return nil, agentSubmitErrorf("invalid resubmit request: %w", err) + } + req.JobName = strings.TrimSpace(req.JobName) + normalizeAgentSubmitOptionalString(&req.Name) + normalizeAgentSubmitOptionalString(&req.CPU) + normalizeAgentSubmitOptionalString(&req.Memory) + normalizeAgentSubmitOptionalString(&req.GPUModel) + if req.JobName == "" { + return nil, agentSubmitErrorf("job_name is required") + } + if token.Username == "" { + return nil, agentSubmitErrorf("user identity is unavailable for resubmit") + } + + jobRecord, err := getJob(ctx, req.JobName, &token) + if err != nil { + return nil, agentSubmitErrorf("job not found or access denied: %w", err) + } + sourceJob := jobRecord.Attributes.Data() + if sourceJob == nil { + return nil, agentSubmitErrorf("job spec is unavailable for resubmit") + } + + clonedJob := sourceJob.DeepCopy() + appliedOverrides, err := applyAgentSubmitResubmitOverrides(clonedJob, req.CPU, req.Memory, req.GPUCount, req.GPUModel) + if err != nil { + return nil, err + } + + prefix := agentSubmitJobNamePrefix(jobRecord.JobName) + newJobName := utils.GenerateJobName(prefix, token.Username) + baseURL := agentSubmitBaseURLFromJobName(newJobName) + + clonedJob.ObjectMeta = metav1.ObjectMeta{ + Name: newJobName, + Namespace: config.GetConfig().Namespaces.Job, + Labels: copyAgentSubmitStringMap(clonedJob.Labels), + Annotations: copyAgentSubmitStringMap(clonedJob.Annotations), + } + clonedJob.Status = batch.JobStatus{} + clonedJob.ResourceVersion = "" + clonedJob.UID = "" + clonedJob.CreationTimestamp = metav1.Time{} + clonedJob.ManagedFields = nil + clonedJob.OwnerReferences = nil + clonedJob.Finalizers = nil + clonedJob.DeletionTimestamp = nil + + if clonedJob.Labels == nil { + clonedJob.Labels = map[string]string{} + } + clonedJob.Labels[crclient.LabelKeyBaseURL] = baseURL + if clonedJob.Annotations == nil { + clonedJob.Annotations = map[string]string{} + } + if req.Name != nil && strings.TrimSpace(*req.Name) != "" { + clonedJob.Annotations[AnnotationKeyTaskName] = strings.TrimSpace(*req.Name) + appliedOverrides["name"] = strings.TrimSpace(*req.Name) + } else if clonedJob.Annotations[AnnotationKeyTaskName] == "" { + clonedJob.Annotations[AnnotationKeyTaskName] = jobRecord.Name + } + + for idx := range clonedJob.Spec.Tasks { + task := &clonedJob.Spec.Tasks[idx] + task.Template.ResourceVersion = "" + task.Template.UID = "" + task.Template.CreationTimestamp = metav1.Time{} + task.Template.ManagedFields = nil + if task.Template.Labels == nil { + task.Template.Labels = map[string]string{} + } + task.Template.Labels[crclient.LabelKeyBaseURL] = baseURL + task.Template.Labels[crclient.LabelKeyTaskType] = clonedJob.Labels[crclient.LabelKeyTaskType] + task.Template.Labels[crclient.LabelKeyTaskUser] = clonedJob.Labels[crclient.LabelKeyTaskUser] + if accountName := clonedJob.Labels[crclient.LalbeKeyTaskAccount]; accountName != "" { + task.Template.Labels[crclient.LalbeKeyTaskAccount] = accountName + } + } + + if err := s.mgr.client.Create(ctx, clonedJob); err != nil { + return nil, agentSubmitErrorf("failed to create resubmitted job: %w", err) + } + if err := s.ensureResubmittedJobAccess(ctx, clonedJob); err != nil { + return map[string]any{ + "sourceJobName": jobRecord.JobName, + "jobName": newJobName, + "status": "created", + "warning": err.Error(), + }, nil + } + + return map[string]any{ + "sourceJobName": jobRecord.JobName, + "jobName": newJobName, + "displayName": clonedJob.Annotations[AnnotationKeyTaskName], + "status": "created", + "overrides": appliedOverrides, + }, nil +} + +func (s *agentJobSubmitter) SubmitJupyterJob( + ctx context.Context, + token util.JWTMessage, + rawReq json.RawMessage, +) (any, error) { + var req CreateJupyterReq + if err := json.Unmarshal(rawReq, &req); err != nil { + return nil, agentSubmitErrorf("invalid jupyter request: %w", err) + } + + // Resolve scheduling metadata the same way the /v1/vcjobs/jupyter HTTP + // handler does, so agent-submitted jobs carry the same prequeue / + // backfill / tolerance annotations as user-submitted ones. + scheduleType, err := req.validateScheduleOptions(true) + if err != nil { + return nil, agentSubmitErrorf("invalid schedule options: %w", err) + } + scheduleMetadata, err := s.mgr.resolveJobScheduleMetadata(ctx, scheduleType) + if err != nil { + return nil, agentSubmitErrorf("failed to resolve schedule metadata: %w", err) + } + + if err := s.prepareJobCreate(ctx, token, scheduleType, true, req.AlertEnabled); err != nil { + return nil, err + } + + jobName := utils.GenerateJobName("jpt", token.Username) + baseURL := jobName[4:] + + jupyterCommand := fmt.Sprintf( + "jupyter lab --ip=0.0.0.0 --no-browser --allow-root "+ + "--notebook-dir=/home/%s --NotebookApp.base_url=/ingress/%s/ "+ + "--ResourceUseDisplay.track_cpu_percent=True", + token.Username, baseURL) + + commandArgs := []string{ + "/bin/bash", + "-c", + fmt.Sprintf("/usr/local/bin/unified-start.sh %s", jupyterCommand), + } + + labels, jobAnnotations, podAnnotations := getLabelAndAnnotations( + CraterJobTypeJupyter, token, baseURL, &req.CreateJobCommon, scheduleMetadata, + ) + + podSpec, err := generateInteractivePodSpec( + ctx, token, &req.CreateJobCommon, req.Resource, req.Image, + commandArgs, string(CraterJobTypeJupyter), req.CpuPinningEnabled, + ) + if err != nil { + return nil, err + } + + queueName := token.AccountName + if token.AccountID != model.DefaultAccountID { + queueName = vcqueue.GetUserQueueName(token.AccountID, token.UserID) + } + + job := buildInteractiveVolcanoJob(jobName, labels, jobAnnotations, podAnnotations, &podSpec, queueName) + + if err := s.mgr.client.Create(ctx, &job); err != nil { + return nil, err + } + + port := &v1.ServicePort{ + Name: "notebook", + Port: JupyterPort, + TargetPort: intstr.FromInt(JupyterPort), + Protocol: v1.ProtocolTCP, + } + + ingressPath, err := s.mgr.serviceManager.CreateIngressWithPrefix( + ctx, + []metav1.OwnerReference{ + *metav1.NewControllerRef(&job, batch.SchemeGroupVersion.WithKind("Job")), + }, + labels, port, config.GetConfig().Host, baseURL, + ) + if err != nil { + return nil, agentSubmitErrorf("failed to create ingress: %v", err) + } + log.Printf("Ingress created at path: %s", ingressPath) + + if err := s.mgr.CreateForwardIngresses(ctx, &job, req.Forwards, labels, token.Username); err != nil { + return nil, err + } + return &job, nil +} + +func (s *agentJobSubmitter) SubmitWebIDEJob( + ctx context.Context, + token util.JWTMessage, + rawReq json.RawMessage, +) (any, error) { + var req CreateJupyterReq + if err := json.Unmarshal(rawReq, &req); err != nil { + return nil, agentSubmitErrorf("invalid webide request: %w", err) + } + + scheduleType, err := req.validateScheduleOptions(true) + if err != nil { + return nil, agentSubmitErrorf("invalid schedule options: %w", err) + } + scheduleMetadata, err := s.mgr.resolveJobScheduleMetadata(ctx, scheduleType) + if err != nil { + return nil, agentSubmitErrorf("failed to resolve schedule metadata: %w", err) + } + + if err := s.prepareJobCreate(ctx, token, scheduleType, true, req.AlertEnabled); err != nil { + return nil, err + } + + jobName := utils.GenerateJobName("vsc", token.Username) + baseURL := jobName[4:] + webIDECommand := fmt.Sprintf("code-server --bind-addr 0.0.0.0:%d", JupyterPort) + commandArgs := []string{ + "/bin/bash", + "-c", + fmt.Sprintf("/usr/local/bin/unified-start.sh %s", webIDECommand), + } + + labels, jobAnnotations, podAnnotations := getLabelAndAnnotations( + CraterJobTypeWebIDE, token, baseURL, &req.CreateJobCommon, scheduleMetadata, + ) + + podSpec, err := generateInteractivePodSpec( + ctx, token, &req.CreateJobCommon, req.Resource, req.Image, + commandArgs, string(CraterJobTypeWebIDE), req.CpuPinningEnabled, + ) + if err != nil { + return nil, err + } + + job := batch.Job{ + ObjectMeta: metav1.ObjectMeta{ + Name: jobName, + Namespace: config.GetConfig().Namespaces.Job, + Labels: labels, + Annotations: jobAnnotations, + }, + Spec: batch.JobSpec{ + TTLSecondsAfterFinished: ptr.To(utils.ThreeDaySeconds), + MinAvailable: 1, + MaxRetry: 1, + Plugins: volcanoPlugins, + SchedulerName: VolcanoSchedulerName, + Queue: vcqueue.ResolveJobQueueName(token), + Policies: []batch.LifecyclePolicy{ + {Action: bus.RestartJobAction, Event: bus.PodEvictedEvent}, + }, + Tasks: []batch.TaskSpec{ + { + Replicas: 1, + Template: v1.PodTemplateSpec{ + ObjectMeta: metav1.ObjectMeta{ + Labels: labels, + Annotations: podAnnotations, + }, + Spec: podSpec, + }, + }, + }, + }, + } + + if err := s.mgr.submitJob(ctx, token, &job); err != nil { + return nil, err + } + return &job, nil +} + +func (s *agentJobSubmitter) SubmitTrainingJob( + ctx context.Context, + token util.JWTMessage, + rawReq json.RawMessage, +) (any, error) { + var req CreateCustomReq + if err := json.Unmarshal(rawReq, &req); err != nil { + return nil, agentSubmitErrorf("invalid training request: %w", err) + } + + // Match /v1/vcjobs/custom: resolve schedule metadata before quota checks. + scheduleType, err := req.validateScheduleOptions(true) + if err != nil { + return nil, agentSubmitErrorf("invalid schedule options: %w", err) + } + scheduleMetadata, err := s.mgr.resolveJobScheduleMetadata(ctx, scheduleType) + if err != nil { + return nil, agentSubmitErrorf("failed to resolve schedule metadata: %w", err) + } + + if err := s.prepareJobCreate(ctx, token, scheduleType, false, req.AlertEnabled); err != nil { + return nil, err + } + + jobName := utils.GenerateJobName("sg", token.Username) + baseURL := jobName[3:] + + labels, jobAnnotations, podAnnotations := getLabelAndAnnotations( + CraterJobTypeCustom, token, baseURL, &req.CreateJobCommon, scheduleMetadata, + ) + + podSpec, err := GenerateCustomPodSpec(ctx, token, &req) + if err != nil { + return nil, err + } + + queueName := token.AccountName + if token.AccountID != model.DefaultAccountID { + queueName = vcqueue.GetUserQueueName(token.AccountID, token.UserID) + } + + job := buildTrainingVolcanoJob(jobName, labels, jobAnnotations, podAnnotations, &podSpec, queueName) + + if err := s.mgr.client.Create(ctx, &job); err != nil { + return nil, err + } + + if err := s.mgr.CreateForwardIngresses(ctx, &job, req.Forwards, labels, token.Username); err != nil { + return nil, err + } + return &job, nil +} + +func (s *agentJobSubmitter) SubmitTensorflowJob( + ctx context.Context, + token util.JWTMessage, + rawReq json.RawMessage, +) (any, error) { + var req CreateTensorflowReq + if err := json.Unmarshal(rawReq, &req); err != nil { + return nil, agentSubmitErrorf("invalid tensorflow request: %w", err) + } + + scheduleType, err := req.validateScheduleOptions(false) + if err != nil { + return nil, agentSubmitErrorf("invalid schedule options: %w", err) + } + scheduleMetadata, err := s.mgr.resolveJobScheduleMetadata(ctx, scheduleType) + if err != nil { + return nil, agentSubmitErrorf("failed to resolve schedule metadata: %w", err) + } + + jobResources := utils.CalculateReplicatedResources( + req.Tasks, + func(task TaskReq) v1.ResourceList { + return task.Resource + }, + func(task TaskReq) int32 { + return task.Replicas + }, + ) + + if err := s.prepareJobCreate(ctx, token, scheduleType, false, req.AlertEnabled); err != nil { + return nil, err + } + + jobName := utils.GenerateJobName("tf", token.Username) + baseURL := jobName[3:] + + volumes, volumeMounts, err := GenerateVolumeMounts(ctx, req.VolumeMounts, token) + if err != nil { + return nil, err + } + baseAffinity := GenerateNodeAffinity(req.Selectors, jobResources) + baseTolerations := GenerateTaintTolerationsForAccount(token) + envs := GenerateEnvs(ctx, token, req.Envs) + + labels, jobAnnotations, podAnnotations := getLabelAndAnnotations( + CraterJobTypeTensorflow, + token, + baseURL, + &req.CreateJobCommon, + scheduleMetadata, + ) + + tasks, minAvailable := buildAgentParallelTaskSpecs( + &req, baseAffinity, baseTolerations, volumes, volumeMounts, envs, labels, podAnnotations, CraterJobTypeTensorflow, + ) + + job := batch.Job{ + ObjectMeta: metav1.ObjectMeta{ + Name: jobName, + Namespace: config.GetConfig().Namespaces.Job, + Labels: labels, + Annotations: jobAnnotations, + }, + Spec: batch.JobSpec{ + TTLSecondsAfterFinished: ptr.To(utils.SevenDaySeconds), + MinAvailable: minAvailable, + SchedulerName: VolcanoSchedulerName, + Plugins: map[string][]string{ + "env": {}, + "svc": {}, + }, + Policies: []batch.LifecyclePolicy{ + { + Action: bus.RestartJobAction, + Event: bus.PodEvictedEvent, + }, + }, + Queue: vcqueue.ResolveJobQueueName(token), + Tasks: tasks, + }, + } + + if err := s.mgr.submitJob(ctx, token, &job); err != nil { + return nil, err + } + return &job, nil +} + +func (s *agentJobSubmitter) SubmitPytorchJob( + ctx context.Context, + token util.JWTMessage, + rawReq json.RawMessage, +) (any, error) { + var req CreateTensorflowReq + if err := json.Unmarshal(rawReq, &req); err != nil { + return nil, agentSubmitErrorf("invalid pytorch request: %w", err) + } + + scheduleType, err := req.validateScheduleOptions(false) + if err != nil { + return nil, agentSubmitErrorf("invalid schedule options: %w", err) + } + scheduleMetadata, err := s.mgr.resolveJobScheduleMetadata(ctx, scheduleType) + if err != nil { + return nil, agentSubmitErrorf("failed to resolve schedule metadata: %w", err) + } + + jobResources := utils.CalculateReplicatedResources( + req.Tasks, + func(task TaskReq) v1.ResourceList { + return task.Resource + }, + func(task TaskReq) int32 { + return task.Replicas + }, + ) + + if err := s.prepareJobCreate(ctx, token, scheduleType, false, req.AlertEnabled); err != nil { + return nil, err + } + + jobName := utils.GenerateJobName("pyt", token.Username) + baseURL := jobName[4:] + + volumes, volumeMounts, err := GenerateVolumeMounts(ctx, req.VolumeMounts, token) + if err != nil { + return nil, err + } + baseAffinity := GenerateNodeAffinity(req.Selectors, jobResources) + baseTolerations := GenerateTaintTolerationsForAccount(token) + envs := GenerateEnvs(ctx, token, req.Envs) + + labels, jobAnnotations, podAnnotations := getLabelAndAnnotations( + CraterJobTypePytorch, + token, + baseURL, + &req.CreateJobCommon, + scheduleMetadata, + ) + + tasks, minAvailable := buildAgentParallelTaskSpecs( + &req, baseAffinity, baseTolerations, volumes, volumeMounts, envs, labels, podAnnotations, CraterJobTypePytorch, + ) + + job := batch.Job{ + ObjectMeta: metav1.ObjectMeta{ + Name: jobName, + Namespace: config.GetConfig().Namespaces.Job, + Labels: labels, + Annotations: jobAnnotations, + }, + Spec: batch.JobSpec{ + TTLSecondsAfterFinished: ptr.To(utils.SevenDaySeconds), + MinAvailable: minAvailable, + SchedulerName: VolcanoSchedulerName, + Plugins: map[string][]string{ + string(CraterJobTypePytorch): {pytorchPluginMasterArg, pytorchPluginWorkerArg, pytorchPluginPortArg}, + }, + Policies: []batch.LifecyclePolicy{ + { + Action: bus.RestartJobAction, + Event: bus.PodEvictedEvent, + }, + }, + Queue: vcqueue.ResolveJobQueueName(token), + Tasks: tasks, + }, + } + + if err := s.mgr.submitJob(ctx, token, &job); err != nil { + return nil, err + } + return &job, nil +} + +func normalizeAgentSubmitOptionalString(value **string) { + if value == nil || *value == nil { + return + } + trimmed := strings.TrimSpace(**value) + if trimmed == "" { + *value = nil + return + } + *value = &trimmed +} + +func agentSubmitJobNamePrefix(jobName string) string { + parts := strings.SplitN(jobName, "-", 2) + if len(parts) > 0 && parts[0] != "" { + return parts[0] + } + return "job" +} + +func agentSubmitBaseURLFromJobName(jobName string) string { + parts := strings.SplitN(jobName, "-", 2) + if len(parts) == 2 && parts[1] != "" { + return parts[1] + } + return jobName +} + +func copyAgentSubmitStringMap(src map[string]string) map[string]string { + if len(src) == 0 { + return map[string]string{} + } + dst := make(map[string]string, len(src)) + for k, v := range src { + dst[k] = v + } + return dst +} + +func (s *agentJobSubmitter) ensureResubmittedJobAccess(ctx context.Context, job *batch.Job) error { + if job == nil || s.mgr.serviceManager == nil || len(job.Spec.Tasks) == 0 { + return nil + } + + labels := copyAgentSubmitStringMap(job.Labels) + taskType := labels[crclient.LabelKeyTaskType] + baseURL := labels[crclient.LabelKeyBaseURL] + ownerRefs := []metav1.OwnerReference{ + *metav1.NewControllerRef(job, batch.SchemeGroupVersion.WithKind("Job")), + } + + switch taskType { + case string(model.JobTypeJupyter): + _, err := s.mgr.serviceManager.CreateIngressWithPrefix( + ctx, + ownerRefs, + labels, + &v1.ServicePort{ + Name: "notebook", + Port: JupyterPort, + TargetPort: intstr.FromInt(JupyterPort), + Protocol: v1.ProtocolTCP, + }, + config.GetConfig().Host, + baseURL, + ) + return err + case string(model.JobTypeWebIDE): + username := labels[crclient.LabelKeyTaskUser] + randomPrefix := uuid.New().String()[:5] + _, err := s.mgr.serviceManager.CreateNamedIngress( + ctx, + ownerRefs, + labels, + &v1.ServicePort{ + Name: "webide", + Port: JupyterPort, + TargetPort: intstr.FromInt(JupyterPort), + Protocol: v1.ProtocolTCP, + }, + config.GetConfig().Host, + username, + randomPrefix, + ) + return err + default: + return nil + } +} + +//nolint:gocyclo // Resubmit overrides walk task/container resource specs and preserve GPU model semantics. +func applyAgentSubmitResubmitOverrides( + job *batch.Job, + cpu *string, + memory *string, + gpuCount *int, + gpuModel *string, +) (map[string]any, error) { + if job == nil { + return nil, agentSubmitErrorf("job spec is unavailable for override") + } + + applied := make(map[string]any) + for taskIdx := range job.Spec.Tasks { + task := &job.Spec.Tasks[taskIdx] + for containerIdx := range task.Template.Spec.Containers { + container := &task.Template.Spec.Containers[containerIdx] + if cpu != nil { + quantity, err := resource.ParseQuantity(strings.TrimSpace(*cpu)) + if err != nil { + return nil, agentSubmitErrorf("invalid cpu override: %w", err) + } + if container.Resources.Requests == nil { + container.Resources.Requests = v1.ResourceList{} + } + if container.Resources.Limits == nil { + container.Resources.Limits = v1.ResourceList{} + } + container.Resources.Requests[v1.ResourceCPU] = quantity + container.Resources.Limits[v1.ResourceCPU] = quantity + applied["cpu"] = quantity.String() + } + if memory != nil { + quantity, err := resource.ParseQuantity(strings.TrimSpace(*memory)) + if err != nil { + return nil, agentSubmitErrorf("invalid memory override: %w", err) + } + if container.Resources.Requests == nil { + container.Resources.Requests = v1.ResourceList{} + } + if container.Resources.Limits == nil { + container.Resources.Limits = v1.ResourceList{} + } + container.Resources.Requests[v1.ResourceMemory] = quantity + container.Resources.Limits[v1.ResourceMemory] = quantity + applied["memory"] = quantity.String() + } + + gpuResourceName, changed, err := overrideAgentSubmitGPUResourceRequirements( + &container.Resources, + gpuCount, + gpuModel, + ) + if err != nil { + return nil, err + } + if changed { + if gpuCount != nil { + applied["gpu_count"] = *gpuCount + } + if gpuResourceName != "" { + applied["gpu_resource_name"] = gpuResourceName + if gpuModel != nil && strings.TrimSpace(*gpuModel) != "" { + applied["gpu_model"] = normalizeAgentSubmitGPUModelName(*gpuModel) + } + } + } + } + } + + if len(applied) == 0 { + applied["inherit"] = "original_spec" + } + return applied, nil +} + +//nolint:gocyclo // GPU resource overrides must handle count, model rename, zeroing and inherited resource names. +func overrideAgentSubmitGPUResourceRequirements( + requirements *v1.ResourceRequirements, + gpuCount *int, + gpuModel *string, +) (gpuResourceName string, changed bool, err error) { + if requirements == nil { + return "", false, nil + } + currentGPUKey := detectAgentSubmitGPUResourceName(requirements.Requests) + if currentGPUKey == "" { + currentGPUKey = detectAgentSubmitGPUResourceName(requirements.Limits) + } + if currentGPUKey == "" && gpuCount == nil && gpuModel == nil { + return "", false, nil + } + + targetGPUKey := currentGPUKey + if gpuModel != nil && strings.TrimSpace(*gpuModel) != "" { + targetGPUKey = normalizeAgentSubmitGPUResourceName(currentGPUKey, *gpuModel) + } + if targetGPUKey == "" && gpuCount != nil && *gpuCount > 0 { + targetGPUKey = normalizeAgentSubmitGPUResourceName(currentGPUKey, "gpu") + } + if targetGPUKey == "" { + return "", false, nil + } + + changed = false + if requirements.Requests == nil { + requirements.Requests = v1.ResourceList{} + } + if requirements.Limits == nil { + requirements.Limits = v1.ResourceList{} + } + if currentGPUKey != "" && currentGPUKey != targetGPUKey { + moveAgentSubmitResourceQuantity(requirements.Requests, currentGPUKey, targetGPUKey) + moveAgentSubmitResourceQuantity(requirements.Limits, currentGPUKey, targetGPUKey) + changed = true + } + if removeAgentSubmitGPUResourcesExcept(requirements.Requests, targetGPUKey) { + changed = true + } + if removeAgentSubmitGPUResourcesExcept(requirements.Limits, targetGPUKey) { + changed = true + } + + if gpuCount != nil { + if *gpuCount < 0 { + return "", false, agentSubmitErrorf("gpu_count must be non-negative") + } + if *gpuCount == 0 { + if removeAgentSubmitGPUResourcesExcept(requirements.Requests, "") { + changed = true + } + if removeAgentSubmitGPUResourcesExcept(requirements.Limits, "") { + changed = true + } + return "", changed, nil + } + quantity := *resource.NewQuantity(int64(*gpuCount), resource.DecimalSI) + requirements.Requests[targetGPUKey] = quantity + requirements.Limits[targetGPUKey] = quantity + changed = true + } + + return string(targetGPUKey), changed, nil +} + +func detectAgentSubmitGPUResourceName(resources v1.ResourceList) v1.ResourceName { + for name := range resources { + if isAgentSubmitGPUResourceName(string(name)) { + return name + } + } + return "" +} + +func normalizeAgentSubmitGPUModelName(input string) string { + gpuModelName := strings.TrimSpace(strings.ToLower(input)) + return strings.ReplaceAll(gpuModelName, " ", "-") +} + +func normalizeAgentSubmitGPUResourceName(current v1.ResourceName, gpuModel string) v1.ResourceName { + gpuModelName := normalizeAgentSubmitGPUModelName(gpuModel) + if gpuModelName == "" { + return current + } + if strings.Contains(gpuModelName, "/") { + return v1.ResourceName(gpuModelName) + } + vendor := "nvidia.com" + if current != "" { + parts := strings.SplitN(string(current), "/", 2) + if len(parts) == 2 && parts[0] != "" { + vendor = parts[0] + } + } + return v1.ResourceName(fmt.Sprintf("%s/%s", vendor, gpuModelName)) +} + +func moveAgentSubmitResourceQuantity(resources v1.ResourceList, oldName, newName v1.ResourceName) { + if resources == nil || oldName == "" || oldName == newName { + return + } + if quantity, ok := resources[oldName]; ok { + resources[newName] = quantity + delete(resources, oldName) + } +} + +func removeAgentSubmitGPUResourcesExcept(resources v1.ResourceList, keep v1.ResourceName) bool { + if resources == nil { + return false + } + changed := false + for name := range resources { + if name == keep || !isAgentSubmitGPUResourceName(string(name)) { + continue + } + delete(resources, name) + changed = true + } + return changed +} + +func isAgentSubmitGPUResourceName(name string) bool { + normalized := strings.ToLower(strings.TrimSpace(name)) + if normalized == "" { + return false + } + if strings.HasPrefix(normalized, "nvidia.com/") || strings.Contains(normalized, "/gpu") || strings.Contains(normalized, "gpu") { + return true + } + for _, gpuModelName := range []string{"v100", "a100", "h100", "l40s", "rtx4090"} { + if normalized == gpuModelName || strings.HasSuffix(normalized, "/"+gpuModelName) { + return true + } + } + return false +} diff --git a/backend/internal/handler/vcjob/custom.go b/backend/internal/handler/vcjob/custom.go index 7635bbb1f..9ed85e799 100644 --- a/backend/internal/handler/vcjob/custom.go +++ b/backend/internal/handler/vcjob/custom.go @@ -6,10 +6,7 @@ import ( "github.com/gin-gonic/gin" v1 "k8s.io/api/core/v1" - metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/utils/ptr" - batch "volcano.sh/apis/pkg/apis/batch/v1alpha1" - bus "volcano.sh/apis/pkg/apis/bus/v1alpha1" "github.com/raids-lab/crater/internal/resputil" "github.com/raids-lab/crater/internal/util" @@ -102,46 +99,7 @@ func (mgr *VolcanojobMgr) CreateTrainingJob(c *gin.Context) { queueName := vcqueue.ResolveJobQueueName(token) // 6. Create volcano job - job := batch.Job{ - ObjectMeta: metav1.ObjectMeta{ - Name: jobName, - Namespace: config.GetConfig().Namespaces.Job, - Labels: labels, - Annotations: jobAnnotations, - }, - Spec: batch.JobSpec{ - TTLSecondsAfterFinished: ptr.To(utils.SevenDaySeconds), - MinAvailable: 1, - MaxRetry: 1, - SchedulerName: VolcanoSchedulerName, - Queue: queueName, - Plugins: volcanoPlugins, - Policies: []batch.LifecyclePolicy{ - { - Action: bus.RestartJobAction, - Event: bus.PodEvictedEvent, - }, - }, - Tasks: []batch.TaskSpec{ - { - Replicas: 1, - Template: v1.PodTemplateSpec{ - ObjectMeta: metav1.ObjectMeta{ - Labels: labels, - Annotations: podAnnotations, - }, - Spec: podSpec, - }, - Policies: []batch.LifecyclePolicy{ - { - Action: bus.CompleteJobAction, - Event: bus.TaskCompletedEvent, - }, - }, - }, - }, - }, - } + job := buildTrainingVolcanoJob(jobName, labels, jobAnnotations, podAnnotations, &podSpec, queueName) if err = mgr.submitJob(c, token, &job); err != nil { resputil.Error(c, err.Error(), resputil.NotSpecified) diff --git a/backend/internal/handler/vcjob/jupyter.go b/backend/internal/handler/vcjob/jupyter.go index 0a76f2ebb..796a9838c 100644 --- a/backend/internal/handler/vcjob/jupyter.go +++ b/backend/internal/handler/vcjob/jupyter.go @@ -6,11 +6,8 @@ import ( "github.com/gin-gonic/gin" v1 "k8s.io/api/core/v1" - metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" - "k8s.io/utils/ptr" "sigs.k8s.io/controller-runtime/pkg/client" batch "volcano.sh/apis/pkg/apis/batch/v1alpha1" - bus "volcano.sh/apis/pkg/apis/bus/v1alpha1" "github.com/raids-lab/crater/dao/model" "github.com/raids-lab/crater/internal/resputil" @@ -48,8 +45,6 @@ type ( // @Failure 400 {object} resputil.Response[any] "Request parameter error" // @Failure 500 {object} resputil.Response[any] "Other errors" // @Router /v1/vcjobs/jupyter [post] -// -//nolint:dupl //TODO: refactor similar code with CreateWebIDEJob func (mgr *VolcanojobMgr) CreateJupyterJob(c *gin.Context) { token := util.GetToken(c) @@ -122,7 +117,6 @@ func (mgr *VolcanojobMgr) CreateJupyterJob(c *gin.Context) { req.Resource, req.Image, commandArgs, - JupyterPort, string(CraterJobTypeJupyter), req.CpuPinningEnabled, ) @@ -133,41 +127,7 @@ func (mgr *VolcanojobMgr) CreateJupyterJob(c *gin.Context) { queueName := vcqueue.ResolveJobQueueName(token) // 6. Create volcano job - job := batch.Job{ - ObjectMeta: metav1.ObjectMeta{ - Name: jobName, - Namespace: config.GetConfig().Namespaces.Job, - Labels: labels, - Annotations: jobAnnotations, - }, - Spec: batch.JobSpec{ - // 3 days - TTLSecondsAfterFinished: ptr.To(utils.ThreeDaySeconds), - MinAvailable: 1, - MaxRetry: 1, - Plugins: volcanoPlugins, - SchedulerName: VolcanoSchedulerName, - Queue: queueName, - Policies: []batch.LifecyclePolicy{ - { - Action: bus.RestartJobAction, - Event: bus.PodEvictedEvent, - }, - }, - Tasks: []batch.TaskSpec{ - { - Replicas: 1, - Template: v1.PodTemplateSpec{ - ObjectMeta: metav1.ObjectMeta{ - Labels: labels, - Annotations: podAnnotations, - }, - Spec: podSpec, - }, - }, - }, - }, - } + job := buildInteractiveVolcanoJob(jobName, labels, jobAnnotations, podAnnotations, &podSpec, queueName) if err = mgr.submitJob(c, token, &job); err != nil { resputil.Error(c, err.Error(), resputil.NotSpecified) diff --git a/backend/internal/handler/vcjob/pytorch.go b/backend/internal/handler/vcjob/pytorch.go index b5ba5b523..dd39d6029 100644 --- a/backend/internal/handler/vcjob/pytorch.go +++ b/backend/internal/handler/vcjob/pytorch.go @@ -135,7 +135,7 @@ func (mgr *VolcanojobMgr) CreatePytorchJob(c *gin.Context) { } switch task.Name { - case "master": + case volcanoTaskMaster: taskSpec.Policies = []batch.LifecyclePolicy{ { Action: bus.CompleteJobAction, @@ -146,7 +146,7 @@ func (mgr *VolcanojobMgr) CreatePytorchJob(c *gin.Context) { Event: bus.PodFailedEvent, }, } - case "worker": + case volcanoTaskWorker: taskSpec.Template.Spec.RestartPolicy = v1.RestartPolicyOnFailure } @@ -168,7 +168,7 @@ func (mgr *VolcanojobMgr) CreatePytorchJob(c *gin.Context) { MinAvailable: minAvailable, SchedulerName: VolcanoSchedulerName, Plugins: map[string][]string{ - "pytorch": {"--master=master", "--worker=worker", "--port=23456"}, + string(CraterJobTypePytorch): {pytorchPluginMasterArg, pytorchPluginWorkerArg, pytorchPluginPortArg}, }, Policies: []batch.LifecyclePolicy{ { diff --git a/backend/internal/handler/vcjob/tensorflow.go b/backend/internal/handler/vcjob/tensorflow.go index 0e86c96cb..108d050c9 100644 --- a/backend/internal/handler/vcjob/tensorflow.go +++ b/backend/internal/handler/vcjob/tensorflow.go @@ -215,7 +215,7 @@ func buildTensorflowTasks( }, } - if task.Name == "worker" { + if task.Name == volcanoTaskWorker { taskSpec.Policies = []batch.LifecyclePolicy{ { Action: bus.CompleteJobAction, diff --git a/backend/internal/handler/vcjob/util.go b/backend/internal/handler/vcjob/util.go index 50c118cb9..86fd70074 100644 --- a/backend/internal/handler/vcjob/util.go +++ b/backend/internal/handler/vcjob/util.go @@ -12,10 +12,12 @@ import ( "gopkg.in/yaml.v3" v1 "k8s.io/api/core/v1" "k8s.io/apimachinery/pkg/api/resource" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/client-go/kubernetes/scheme" "k8s.io/client-go/tools/remotecommand" "k8s.io/utils/ptr" batch "volcano.sh/apis/pkg/apis/batch/v1alpha1" + bus "volcano.sh/apis/pkg/apis/bus/v1alpha1" "github.com/gin-gonic/gin" @@ -25,6 +27,7 @@ import ( "github.com/raids-lab/crater/internal/util" "github.com/raids-lab/crater/pkg/config" "github.com/raids-lab/crater/pkg/crclient" + "github.com/raids-lab/crater/pkg/utils" ) var ( @@ -36,6 +39,79 @@ var ( } ) +func buildSingleTaskVolcanoJob( + jobName string, + labels map[string]string, + jobAnnotations map[string]string, + podAnnotations map[string]string, + podSpec *v1.PodSpec, + queueName string, + ttlSeconds int32, + completeTaskOnSuccess bool, +) batch.Job { + task := batch.TaskSpec{ + Replicas: 1, + Template: v1.PodTemplateSpec{ + ObjectMeta: metav1.ObjectMeta{ + Labels: labels, + Annotations: podAnnotations, + }, + Spec: *podSpec, + }, + } + if completeTaskOnSuccess { + task.Policies = []batch.LifecyclePolicy{ + {Action: bus.CompleteJobAction, Event: bus.TaskCompletedEvent}, + } + } + return batch.Job{ + ObjectMeta: metav1.ObjectMeta{ + Name: jobName, + Namespace: config.GetConfig().Namespaces.Job, + Labels: labels, + Annotations: jobAnnotations, + }, + Spec: batch.JobSpec{ + TTLSecondsAfterFinished: ptr.To(ttlSeconds), + MinAvailable: 1, + MaxRetry: 1, + Plugins: volcanoPlugins, + SchedulerName: VolcanoSchedulerName, + Queue: queueName, + Policies: []batch.LifecyclePolicy{ + {Action: bus.RestartJobAction, Event: bus.PodEvictedEvent}, + }, + Tasks: []batch.TaskSpec{task}, + }, + } +} + +func buildInteractiveVolcanoJob( + jobName string, + labels map[string]string, + jobAnnotations map[string]string, + podAnnotations map[string]string, + podSpec *v1.PodSpec, + queueName string, +) batch.Job { + return buildSingleTaskVolcanoJob( + jobName, labels, jobAnnotations, podAnnotations, podSpec, queueName, utils.ThreeDaySeconds, false, + ) +} + +func buildTrainingVolcanoJob( + jobName string, + labels map[string]string, + jobAnnotations map[string]string, + podAnnotations map[string]string, + podSpec *v1.PodSpec, + queueName string, +) batch.Job { + return buildSingleTaskVolcanoJob( + jobName, labels, jobAnnotations, podAnnotations, podSpec, queueName, utils.SevenDaySeconds, true, + ) +} + // buildResourceRequirements creates resource requirements and resize policy based on CPU pinning setting func buildResourceRequirements(resourceList v1.ResourceList, cpuPinningEnabled bool) (v1.ResourceRequirements, []v1.ContainerResizePolicy) { var resourceRequirements v1.ResourceRequirements @@ -603,7 +679,6 @@ func generateInteractivePodSpec( resourceList v1.ResourceList, image ImageBaseInfo, command []string, - port int32, containerName string, cpuPinningEnabled bool, ) (v1.PodSpec, error) { @@ -668,7 +743,7 @@ func generateInteractivePodSpec( Env: envs, Ports: []v1.ContainerPort{ - {ContainerPort: port, Name: containerName, Protocol: v1.ProtocolTCP}, + {ContainerPort: JupyterPort, Name: containerName, Protocol: v1.ProtocolTCP}, }, SecurityContext: &v1.SecurityContext{ RunAsUser: ptr.To(int64(0)), diff --git a/backend/internal/handler/vcjob/vcjob.go b/backend/internal/handler/vcjob/vcjob.go index 10493d14c..bb46f414a 100644 --- a/backend/internal/handler/vcjob/vcjob.go +++ b/backend/internal/handler/vcjob/vcjob.go @@ -137,6 +137,13 @@ func (mgr *VolcanojobMgr) RegisterAdmin(g *gin.RouterGroup) { const ( VolcanoSchedulerName = "volcano" + volcanoTaskMaster = "master" + volcanoTaskWorker = "worker" + + pytorchPluginMasterArg = "--master=master" + pytorchPluginWorkerArg = "--worker=worker" + pytorchPluginPortArg = "--port=23456" + AnnotationKeyUser = "crater.raids.io/user" // 用户名，以小写字母开头 AnnotationKeyTaskName = "crater.raids.io/task-name" // 任务名称（可能是中文） AnnotationKeyTaskTemplate = "crater.raids.io/task-template" // 任务模板 @@ -390,10 +397,10 @@ func (mgr *VolcanojobMgr) getDeleteJobRecord(c *gin.Context, jobName string) (*m return getJob(c, jobName, &token) } -func (mgr *VolcanojobMgr) buildDeleteJobPlan(c *gin.Context, jobRecord *model.Job) (*deleteJobPlan, error) { +func (mgr *VolcanojobMgr) buildDeleteJobPlan(ctx context.Context, jobRecord *model.Job) (*deleteJobPlan, error) { clusterJob := &batch.Job{} namespace := config.GetConfig().Namespaces.Job - if err := mgr.client.Get(c, client.ObjectKey{Name: jobRecord.JobName, Namespace: namespace}, clusterJob); err != nil { + if err := mgr.client.Get(ctx, client.ObjectKey{Name: jobRecord.JobName, Namespace: namespace}, clusterJob); err != nil { if errors.IsNotFound(err) { return &deleteJobPlan{ shouldDeleteRecord: jobRecord.Status != model.Prequeue, @@ -418,13 +425,13 @@ func (mgr *VolcanojobMgr) buildDeleteJobPlan(c *gin.Context, jobRecord *model.Jo }, nil } -func (mgr *VolcanojobMgr) applyDeleteJobPlan(c *gin.Context, record *model.Job, plan *deleteJobPlan) error { +func (mgr *VolcanojobMgr) applyDeleteJobPlan(ctx context.Context, record *model.Job, plan *deleteJobPlan) error { j := query.Job if plan.shouldDeleteRecord { - if err := mgr.settleJobBeforeDelete(c, record, plan.clusterJob); err != nil { + if err := mgr.settleJobBeforeDelete(ctx, record, plan.clusterJob); err != nil { return err } - _, err := j.WithContext(c).Where(j.JobName.Eq(record.JobName)).Delete() + _, err := j.WithContext(ctx).Where(j.JobName.Eq(record.JobName)).Delete() return err } @@ -433,24 +440,24 @@ func (mgr *VolcanojobMgr) applyDeleteJobPlan(c *gin.Context, record *model.Job, finalJob := *record finalJob.Status = model.Deleted finalJob.CompletedTimestamp = completedAt - if err := mgr.billingService.OnJobFinishedSettlement(c.Request.Context(), &finalJob); err != nil { + if err := mgr.billingService.OnJobFinishedSettlement(ctx, &finalJob); err != nil { return err } } - _, err := j.WithContext(c).Where(j.JobName.Eq(record.JobName)).Updates(model.Job{ + _, err := j.WithContext(ctx).Where(j.JobName.Eq(record.JobName)).Updates(model.Job{ Status: model.Deleted, CompletedTimestamp: completedAt, }) return err } -func (mgr *VolcanojobMgr) deleteClusterJob(c *gin.Context, plan *deleteJobPlan) error { +func (mgr *VolcanojobMgr) deleteClusterJob(ctx context.Context, plan *deleteJobPlan) error { if !plan.shouldDeleteJob || plan.clusterJob == nil { return nil } - return mgr.client.Delete(c, plan.clusterJob) + return mgr.client.Delete(ctx, plan.clusterJob) } func (mgr *VolcanojobMgr) deleteJob(c *gin.Context, recordAdminOperation bool) { @@ -473,20 +480,20 @@ func (mgr *VolcanojobMgr) deleteJob(c *gin.Context, recordAdminOperation bool) { return } - plan, err := mgr.buildDeleteJobPlan(c, jobRecord) + plan, err := mgr.buildDeleteJobPlan(c.Request.Context(), jobRecord) if err != nil { recordDeleteJobOperation(constants.OpStatusFailed, err.Error(), jobRecord, nil) resputil.Error(c, err.Error(), resputil.NotSpecified) return } - if err := mgr.applyDeleteJobPlan(c, jobRecord, plan); err != nil { + if err := mgr.applyDeleteJobPlan(c.Request.Context(), jobRecord, plan); err != nil { recordDeleteJobOperation(constants.OpStatusFailed, err.Error(), jobRecord, nil) resputil.Error(c, err.Error(), resputil.NotSpecified) return } - if err := mgr.deleteClusterJob(c, plan); err != nil { + if err := mgr.deleteClusterJob(c.Request.Context(), plan); err != nil { recordDeleteJobOperation(constants.OpStatusFailed, err.Error(), jobRecord, nil) resputil.Error(c, err.Error(), resputil.NotSpecified) return @@ -510,14 +517,14 @@ func (mgr *VolcanojobMgr) notifyDeletedPrequeue(shouldDeleteRecord bool) { mgr.prequeueWatcher.RequestFullScan() } -func (mgr *VolcanojobMgr) settleJobBeforeDelete(c *gin.Context, record *model.Job, job *batch.Job) error { +func (mgr *VolcanojobMgr) settleJobBeforeDelete(ctx context.Context, record *model.Job, job *batch.Job) error { if mgr.billingService == nil || record == nil { return nil } finalJob := *record finalJob.CompletedTimestamp = resolveDeleteSettlementTime(record, job) - return mgr.billingService.OnJobFinishedSettlement(c.Request.Context(), &finalJob) + return mgr.billingService.OnJobFinishedSettlement(ctx, &finalJob) } func resolveDeleteSettlementTime(record *model.Job, job *batch.Job) time.Time { diff --git a/backend/internal/handler/vcjob/webide.go b/backend/internal/handler/vcjob/webide.go index 5aa1014ec..7483848b3 100644 --- a/backend/internal/handler/vcjob/webide.go +++ b/backend/internal/handler/vcjob/webide.go @@ -35,8 +35,6 @@ import ( // @Failure 400 {object} resputil.Response[any] "Request parameter error" // @Failure 500 {object} resputil.Response[any] "Other errors" // @Router /v1/vcjobs/webide [post] -// -//nolint:dupl //TODO: refactor similar code with CreateJupyterJob func (mgr *VolcanojobMgr) CreateWebIDEJob(c *gin.Context) { token := util.GetToken(c) @@ -105,7 +103,6 @@ func (mgr *VolcanojobMgr) CreateWebIDEJob(c *gin.Context) { req.Resource, req.Image, commandArgs, - JupyterPort, string(CraterJobTypeWebIDE), req.CpuPinningEnabled, ) diff --git a/backend/internal/register.go b/backend/internal/register.go index 5a1d5a09b..64d511e7f 100644 --- a/backend/internal/register.go +++ b/backend/internal/register.go @@ -4,6 +4,7 @@ import ( "k8s.io/klog/v2" "github.com/raids-lab/crater/internal/handler" + _ "github.com/raids-lab/crater/internal/handler/agent" _ "github.com/raids-lab/crater/internal/handler/aijob" _ "github.com/raids-lab/crater/internal/handler/image" _ "github.com/raids-lab/crater/internal/handler/operations" diff --git a/backend/internal/route.go b/backend/internal/route.go index a63bd1b24..3e5eccb99 100644 --- a/backend/internal/route.go +++ b/backend/internal/route.go @@ -67,4 +67,14 @@ func (b *Backend) RegisterService(conf *handler.RegisterConfig) { for _, mgr := range managers { mgr.RegisterAdmin(adminRouter.Group(mgr.GetName())) } + + ///////////////////////////////////////////////////////////// + //// Internal routers (service-to-service, no user JWT) //// + ///////////////////////////////////////////////////////////// + internalRouter := b.Group("/internal") + for _, mgr := range managers { + if ir, ok := mgr.(handler.InternalRouter); ok { + ir.RegisterInternal(internalRouter.Group(mgr.GetName())) + } + } } diff --git a/backend/internal/service/agent_service.go b/backend/internal/service/agent_service.go new file mode 100644 index 000000000..338db14ee --- /dev/null +++ b/backend/internal/service/agent_service.go @@ -0,0 +1,810 @@ +package service + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "strings" + "time" + + "gorm.io/datatypes" + "gorm.io/gorm" + "k8s.io/klog/v2" + + "github.com/raids-lab/crater/dao/model" + "github.com/raids-lab/crater/dao/query" + "github.com/raids-lab/crater/internal/bizerr" +) + +type toolCallAuditMetadata struct { + ExecutionBackend string +} + +// AgentService encapsulates DB operations for the Agent feature. +type AgentService struct { + db *gorm.DB +} + +const ( + agentFeedbackStatusSubmitted = "submitted" + agentSessionTitleMaxRunes = 100 + agentTurnStatusCompleted = "completed" + agentTurnStatusFailed = "failed" + agentTurnStatusCancelled = "cancelled" //nolint:misspell // Stored status matches existing frontend/API spelling. +) + +var agentToolAuditCompatFields = []string{ + "ExecutionBackend", +} + +var agentToolAuditCompatColumns = []string{ + "execution_backend", +} + +func parseToolCallAuditMeta(_, toolResult json.RawMessage) toolCallAuditMetadata { + meta := toolCallAuditMetadata{} + + resultMap := map[string]any{} + if len(toolResult) > 0 { + _ = json.Unmarshal(toolResult, &resultMap) + } + readMeta := func(source map[string]any) { + if source == nil { + return + } + if v, ok := source["execution_backend"].(string); ok && strings.TrimSpace(v) != "" { + meta.ExecutionBackend = strings.TrimSpace(v) + } + } + readMeta(resultMap) + if nested, ok := resultMap["_audit"].(map[string]any); ok { + readMeta(nested) + } + return meta +} + +func isMissingAgentToolAuditColumnError(err error) bool { + if err == nil { + return false + } + message := strings.ToLower(err.Error()) + if !strings.Contains(message, "agent_tool_calls") || !strings.Contains(message, "does not exist") { + return false + } + for _, column := range agentToolAuditCompatColumns { + if strings.Contains(message, fmt.Sprintf("column %q", column)) { + return true + } + } + return false +} + +func stripUnsupportedAgentToolAuditUpdates(updates map[string]any) map[string]any { + if len(updates) == 0 { + return updates + } + compat := make(map[string]any, len(updates)) + for key, value := range updates { + compat[key] = value + } + for _, column := range agentToolAuditCompatColumns { + delete(compat, column) + } + return compat +} + +func (s *AgentService) createToolCallWithCompat(ctx context.Context, toolCall *model.AgentToolCall) error { + err := s.db.WithContext(ctx).Create(toolCall).Error + if err == nil || !isMissingAgentToolAuditColumnError(err) { + return err + } + return s.db.WithContext(ctx).Omit(agentToolAuditCompatFields...).Create(toolCall).Error +} + +var ErrAgentSessionPinningUnavailable = bizerr.ServiceError.ServiceUnavailable.New( + "agent session pinning is unavailable until database migration completes", +) +var ErrAgentSessionDeleted = bizerr.NotFound.DataBaseNotFound.New("agent session was deleted") + +// NewAgentService creates a new AgentService. +func NewAgentService() *AgentService { + return &AgentService{ + db: query.GetDB(), + } +} + +// CreateSession creates a new agent session in the database. +func (s *AgentService) CreateSession( + ctx context.Context, + sessionID string, + userID, accountID uint, + title string, + pageContext json.RawMessage, + source string, +) (*model.AgentSession, error) { + // Title is varchar(255); truncate long messages to fit. + if runeTitle := []rune(title); len(runeTitle) > agentSessionTitleMaxRunes { + title = string(runeTitle[:agentSessionTitleMaxRunes]) + "..." + } + source = normalizeAgentSessionSource(source) + session := &model.AgentSession{ + SessionID: sessionID, + UserID: userID, + AccountID: accountID, + Title: title, + Source: source, + } + if len(pageContext) > 0 { + session.PageContext = datatypes.JSON(pageContext) + } + if err := s.db.WithContext(ctx).Create(session).Error; err != nil { + return nil, err + } + return session, nil +} + +// GetSession retrieves a session by sessionID. +func (s *AgentService) GetSession(ctx context.Context, sessionID string) (*model.AgentSession, error) { + var session model.AgentSession + if err := s.db.WithContext(ctx).Where("session_id = ?", sessionID).First(&session).Error; err != nil { + return nil, err + } + return &session, nil +} + +func (s *AgentService) GetOwnedSession(ctx context.Context, sessionID string, userID uint) (*model.AgentSession, error) { + session, err := s.GetSession(ctx, sessionID) + if err != nil { + return nil, err + } + if session.UserID != userID { + return nil, bizerr.NotFound.DataBaseNotFound.New("session not found") + } + return session, nil +} + +// GetOrCreateSession retrieves an existing session or creates a new one. +func (s *AgentService) GetOrCreateSession( + ctx context.Context, + sessionID string, + userID, accountID uint, + title string, + pageContext json.RawMessage, +) (*model.AgentSession, bool, error) { + return s.GetOrCreateSessionWithSource(ctx, sessionID, userID, accountID, title, pageContext, agentSessionSourceChat) +} + +func (s *AgentService) GetOrCreateSessionWithSource( + ctx context.Context, + sessionID string, + userID, accountID uint, + title string, + pageContext json.RawMessage, + source string, +) (*model.AgentSession, bool, error) { + var session model.AgentSession + err := s.db.WithContext(ctx).Where("session_id = ?", sessionID).First(&session).Error + if err == nil { + return &session, false, nil + } + if !errors.Is(err, gorm.ErrRecordNotFound) { + return nil, false, err + } + var deletedSession model.AgentSession + deletedErr := s.db.WithContext(ctx).Unscoped().Where("session_id = ?", sessionID).First(&deletedSession).Error + if deletedErr == nil && deletedSession.DeletedAt.Valid { + return nil, false, ErrAgentSessionDeleted + } + if deletedErr != nil && !errors.Is(deletedErr, gorm.ErrRecordNotFound) { + return nil, false, deletedErr + } + created, createErr := s.CreateSession(ctx, sessionID, userID, accountID, title, pageContext, source) + if createErr != nil { + return nil, false, createErr + } + return created, true, nil +} + +// ListSessions returns all chat sessions for a given user, ordered by most recent. +// Only returns source='chat' sessions; ops_audit/system sessions are excluded from the UI. +func (s *AgentService) ListSessions(ctx context.Context, userID uint) ([]*model.AgentSession, error) { + var sessions []*model.AgentSession + sessionQuery := s.db.WithContext(ctx).Where("user_id = ?", userID) + // Exclude non-chat sessions from the UI listing. + if s.db.WithContext(ctx).Migrator().HasColumn(&model.AgentSession{}, "Source") { + sessionQuery = sessionQuery.Where("source = ? OR source IS NULL OR source = ''", agentSessionSourceChat) + } + if s.db.WithContext(ctx).Migrator().HasColumn(&model.AgentSession{}, "PinnedAt") { + sessionQuery = sessionQuery. + Order("CASE WHEN pinned_at IS NULL THEN 1 ELSE 0 END ASC"). + Order("pinned_at DESC") + } + if err := sessionQuery.Order("updated_at DESC").Find(&sessions).Error; err != nil { + return nil, err + } + return sessions, nil +} + +// UpdateSessionTitle updates the title of a session. +// +// Trims whitespace, rejects empty input, truncates to 100 runes (with "…" +// suffix) to match CreateSession's validation. Uses UpdateColumn so the +// session's UpdatedAt timestamp is NOT bumped — a rename is not "new activity" +// and should not reorder the session list. +func (s *AgentService) UpdateSessionTitle(ctx context.Context, sessionID, title string) error { + title = strings.TrimSpace(title) + if title == "" { + return bizerr.BadRequest.MissingParameter.New("title cannot be empty") + } + if runeTitle := []rune(title); len(runeTitle) > agentSessionTitleMaxRunes { + title = string(runeTitle[:agentSessionTitleMaxRunes]) + "…" + } + return s.db.WithContext(ctx). + Model(&model.AgentSession{}). + Where("session_id = ?", sessionID). + UpdateColumn("title", title).Error +} + +func (s *AgentService) UpdateSessionOrchestrationMode( + ctx context.Context, + sessionID string, + orchestrationMode string, +) error { + if orchestrationMode == "" { + return nil + } + return s.db.WithContext(ctx). + Model(&model.AgentSession{}). + Where("session_id = ?", sessionID). + Update("last_orchestration_mode", orchestrationMode).Error +} + +func (s *AgentService) UpdateSessionPinned(ctx context.Context, sessionID string, pinned bool) error { + if !s.db.WithContext(ctx).Migrator().HasColumn(&model.AgentSession{}, "PinnedAt") { + return ErrAgentSessionPinningUnavailable + } + if pinned { + now := time.Now() + return s.db.WithContext(ctx). + Model(&model.AgentSession{}). + Where("session_id = ?", sessionID). + Updates(map[string]any{ + "pinned_at": &now, + "updated_at": time.Now(), + }).Error + } + // Unpin: only clear pinned_at, do NOT touch updated_at so the session + // returns to its original position in the time-sorted list. + return s.db.WithContext(ctx). + Model(&model.AgentSession{}). + Where("session_id = ?", sessionID). + Update("pinned_at", nil).Error +} + +func (s *AgentService) DeleteSession(ctx context.Context, sessionID string) error { + return s.db.WithContext(ctx). + Where("session_id = ?", sessionID). + Delete(&model.AgentSession{}).Error +} + +// IncrementMessageCount increments the message count for a session. +func (s *AgentService) IncrementMessageCount(ctx context.Context, sessionID string) error { + return s.db.WithContext(ctx). + Model(&model.AgentSession{}). + Where("session_id = ?", sessionID). + Updates(map[string]any{ + "message_count": gorm.Expr("message_count + 1"), + "updated_at": time.Now(), + }).Error +} + +// SaveMessage saves a message to the agent_messages table. +func (s *AgentService) SaveMessage(ctx context.Context, msg *model.AgentMessage) error { + if err := s.db.WithContext(ctx).Create(msg).Error; err != nil { + return err + } + // Increment message count in session asynchronously to avoid blocking. + go func() { + if err := s.IncrementMessageCount(context.Background(), msg.SessionID); err != nil { + klog.Warningf("[AgentService] Failed to increment message count for session %s: %v", msg.SessionID, err) + } + }() + return nil +} + +// ListMessages returns all messages for a session ordered by creation time. +func (s *AgentService) ListMessages(ctx context.Context, sessionID string) ([]*model.AgentMessage, error) { + var messages []*model.AgentMessage + if err := s.db.WithContext(ctx). + Where("session_id = ?", sessionID). + Order("created_at ASC"). + Find(&messages).Error; err != nil { + return nil, err + } + return messages, nil +} + +func (s *AgentService) ListToolCalls(ctx context.Context, sessionID string) ([]*model.AgentToolCall, error) { + var toolCalls []*model.AgentToolCall + if err := s.db.WithContext(ctx). + Where("session_id = ?", sessionID). + Order("created_at ASC, id ASC"). + Find(&toolCalls).Error; err != nil { + return nil, err + } + return toolCalls, nil +} + +func (s *AgentService) ListToolCallsByTurn(ctx context.Context, turnID string) ([]*model.AgentToolCall, error) { + var toolCalls []*model.AgentToolCall + if err := s.db.WithContext(ctx). + Where("turn_id = ?", turnID). + Order("created_at ASC, id ASC"). + Find(&toolCalls).Error; err != nil { + return nil, err + } + return toolCalls, nil +} + +// LogToolCall records a tool execution in the agent_tool_calls table. +func (s *AgentService) LogToolCall(ctx context.Context, toolCall *model.AgentToolCall) error { + return s.createToolCallWithCompat(ctx, toolCall) +} + +func (s *AgentService) CreateToolCall(ctx context.Context, toolCall *model.AgentToolCall) (*model.AgentToolCall, error) { + if err := s.createToolCallWithCompat(ctx, toolCall); err != nil { + return nil, err + } + return toolCall, nil +} + +func (s *AgentService) GetToolCallByID(ctx context.Context, id uint) (*model.AgentToolCall, error) { + var toolCall model.AgentToolCall + if err := s.db.WithContext(ctx).First(&toolCall, id).Error; err != nil { + return nil, err + } + return &toolCall, nil +} + +func (s *AgentService) UpdateToolCallOutcome( + ctx context.Context, + id uint, + resultStatus string, + toolResult json.RawMessage, + userConfirmed *bool, +) error { + meta := parseToolCallAuditMeta(nil, toolResult) + updates := map[string]any{ + "result_status": resultStatus, + } + if toolResult != nil { + updates["tool_result"] = datatypes.JSON(toolResult) + } + if userConfirmed != nil { + updates["user_confirmed"] = *userConfirmed + } + if meta.ExecutionBackend != "" { + updates["execution_backend"] = meta.ExecutionBackend + } + updateWithFields := func(fields map[string]any) error { + return s.db.WithContext(ctx). + Model(&model.AgentToolCall{}). + Where("id = ?", id). + Updates(fields).Error + } + err := updateWithFields(updates) + if err == nil || !isMissingAgentToolAuditColumnError(err) { + return err + } + return updateWithFields(stripUnsupportedAgentToolAuditUpdates(updates)) +} + +func (s *AgentService) UpdateToolCallArgs(ctx context.Context, id uint, toolArgs json.RawMessage) error { + return s.db.WithContext(ctx). + Model(&model.AgentToolCall{}). + Where("id = ?", id). + Update("tool_args", datatypes.JSON(toolArgs)).Error +} + +func (s *AgentService) GetToolCallByToolCallID(ctx context.Context, toolCallID string) (*model.AgentToolCall, error) { + var toolCall model.AgentToolCall + if err := s.db.WithContext(ctx).Where("tool_call_id = ?", toolCallID).First(&toolCall).Error; err != nil { + return nil, err + } + return &toolCall, nil +} + +func (s *AgentService) CreateTurn(ctx context.Context, turn *model.AgentTurn) (*model.AgentTurn, error) { + if err := s.db.WithContext(ctx).Create(turn).Error; err != nil { + return nil, err + } + return turn, nil +} + +func (s *AgentService) GetTurn(ctx context.Context, turnID string) (*model.AgentTurn, error) { + var turn model.AgentTurn + if err := s.db.WithContext(ctx).Where("turn_id = ?", turnID).First(&turn).Error; err != nil { + return nil, err + } + return &turn, nil +} + +func (s *AgentService) ListTurns(ctx context.Context, sessionID string) ([]*model.AgentTurn, error) { + var turns []*model.AgentTurn + if err := s.db.WithContext(ctx). + Where("session_id = ?", sessionID). + Order("started_at DESC"). + Find(&turns).Error; err != nil { + return nil, err + } + return turns, nil +} + +func (s *AgentService) UpdateTurnStatus( + ctx context.Context, + turnID string, + status string, + finalMessageID *uint, + metadata json.RawMessage, +) error { + updates := map[string]any{ + "status": status, + "updated_at": time.Now(), + } + now := time.Now() + if status == agentTurnStatusCompleted || status == agentTurnStatusFailed || status == agentTurnStatusCancelled { + updates["ended_at"] = &now + } + if finalMessageID != nil { + updates["final_message_id"] = *finalMessageID + } + if metadata != nil { + updates["metadata"] = datatypes.JSON(metadata) + } + return s.db.WithContext(ctx). + Model(&model.AgentTurn{}). + Where("turn_id = ?", turnID). + Updates(updates).Error +} + +func (s *AgentService) CreateRunEvent(ctx context.Context, event *model.AgentRunEvent) (*model.AgentRunEvent, error) { + if event.Sequence == 0 { + nextSequence, err := s.NextRunEventSequence(ctx, event.TurnID) + if err != nil { + return nil, err + } + event.Sequence = nextSequence + } + if err := s.db.WithContext(ctx).Create(event).Error; err != nil { + return nil, err + } + return event, nil +} + +func (s *AgentService) NextRunEventSequence(ctx context.Context, turnID string) (int, error) { + var maxSequence int + row := s.db.WithContext(ctx). + Model(&model.AgentRunEvent{}). + Where("turn_id = ?", turnID). + Select("COALESCE(MAX(sequence), 0)"). + Row() + if err := row.Scan(&maxSequence); err != nil { + return 0, err + } + return maxSequence + 1, nil +} + +func (s *AgentService) ListRunEvents(ctx context.Context, turnID string) ([]*model.AgentRunEvent, error) { + var events []*model.AgentRunEvent + if err := s.db.WithContext(ctx). + Where("turn_id = ?", turnID). + Order("sequence ASC, created_at ASC"). + Find(&events).Error; err != nil { + return nil, err + } + return events, nil +} + +// LogToolCallAsync records a tool execution asynchronously to avoid blocking the caller. +func (s *AgentService) LogToolCallAsync( + sessionID, + toolName string, + toolArgs, + toolResult json.RawMessage, + resultStatus string, + latencyMs int, + turnID, + toolCallID, + agentID, + agentRole, + source string, +) { + go func() { + meta := parseToolCallAuditMeta(toolArgs, toolResult) + tc := &model.AgentToolCall{ + SessionID: sessionID, + TurnID: turnID, + ToolCallID: toolCallID, + AgentID: agentID, + AgentRole: agentRole, + Source: normalizeAgentToolCallSource(source), + ToolName: toolName, + ToolArgs: datatypes.JSON(toolArgs), + ToolResult: datatypes.JSON(toolResult), + ResultStatus: resultStatus, + LatencyMs: latencyMs, + CreatedAt: time.Now(), + } + if meta.ExecutionBackend != "" { + tc.ExecutionBackend = meta.ExecutionBackend + } + if err := s.createToolCallWithCompat(context.Background(), tc); err != nil { + klog.Warningf("[AgentService] Failed to log tool call %s for session %s: %v", toolName, sessionID, err) + } + }() +} + +// Feedback. + +var ErrFeedbackAlreadySubmitted = bizerr.Conflict.ResourceStatusError.New( + "feedback already submitted and cannot be modified", +) + +// UpsertFeedback creates or updates a draft feedback. +// Returns the feedback record and a boolean indicating if it was newly created. +func (s *AgentService) UpsertFeedback(ctx context.Context, fb *model.AgentFeedback) (*model.AgentFeedback, bool, error) { + var existing model.AgentFeedback + err := s.db.WithContext(ctx). + Where("user_id = ? AND target_type = ? AND target_id = ?", fb.UserID, fb.TargetType, fb.TargetID). + First(&existing).Error + + if err == nil { + // Record exists — check immutability + if existing.Status == agentFeedbackStatusSubmitted { + return nil, false, ErrFeedbackAlreadySubmitted + } + updates := map[string]any{ + "rating": fb.Rating, + "tags": fb.Tags, + "dimensions": fb.Dimensions, + "comment": fb.Comment, + "updated_at": time.Now(), + } + if err := s.db.WithContext(ctx).Model(&existing).Updates(updates).Error; err != nil { + return nil, false, err + } + existing.Rating = fb.Rating + existing.Tags = fb.Tags + existing.Dimensions = fb.Dimensions + existing.Comment = fb.Comment + return &existing, false, nil + } + + if !errors.Is(err, gorm.ErrRecordNotFound) { + return nil, false, err + } + + // New record + fb.Status = "draft" + if err := s.db.WithContext(ctx).Create(fb).Error; err != nil { + return nil, false, err + } + return fb, true, nil +} + +// SubmitFeedback transitions a draft feedback to submitted (immutable). +func (s *AgentService) SubmitFeedback(ctx context.Context, userID uint, targetType, targetID string) (*model.AgentFeedback, error) { + var fb model.AgentFeedback + err := s.db.WithContext(ctx). + Where("user_id = ? AND target_type = ? AND target_id = ?", userID, targetType, targetID). + First(&fb).Error + if err != nil { + return nil, err + } + if fb.Status == agentFeedbackStatusSubmitted { + return nil, ErrFeedbackAlreadySubmitted + } + now := time.Now() + updates := map[string]any{ + "status": agentFeedbackStatusSubmitted, + "submitted_at": &now, + "updated_at": now, + } + if err := s.db.WithContext(ctx).Model(&fb).Updates(updates).Error; err != nil { + return nil, err + } + fb.Status = agentFeedbackStatusSubmitted + fb.SubmittedAt = &now + return &fb, nil +} + +// ListFeedbacks returns all feedbacks for a user in a given session. +func (s *AgentService) ListFeedbacks(ctx context.Context, sessionID string, userID uint) ([]*model.AgentFeedback, error) { + var feedbacks []*model.AgentFeedback + if err := s.db.WithContext(ctx). + Where("session_id = ? AND user_id = ?", sessionID, userID). + Order("created_at ASC"). + Find(&feedbacks).Error; err != nil { + return nil, err + } + return feedbacks, nil +} + +// FeedbackStats holds aggregated feedback statistics. +type FeedbackStats struct { + Total int64 `json:"total"` + ThumbsUp int64 `json:"thumbsUp"` + ThumbsDown int64 `json:"thumbsDown"` + AvgDimensions map[string]float64 `json:"avgDimensions"` + TopTags []TagCount `json:"topTags"` +} + +// TagCount represents a tag and its occurrence count. +type TagCount struct { + Tag string `json:"tag"` + Count int64 `json:"count"` +} + +// GetFeedbackStats returns aggregated feedback statistics for admins. +func (s *AgentService) GetFeedbackStats(ctx context.Context, from, to *time.Time) (*FeedbackStats, error) { + db := s.db.WithContext(ctx).Model(&model.AgentFeedback{}).Where("status = ?", agentFeedbackStatusSubmitted) + if from != nil { + db = db.Where("submitted_at >= ?", *from) + } + if to != nil { + db = db.Where("submitted_at <= ?", *to) + } + + stats := &FeedbackStats{AvgDimensions: map[string]float64{}} + + // Total / thumbs up / thumbs down + db.Count(&stats.Total) + db.Where("rating = 1").Count(&stats.ThumbsUp) + // Re-scope for thumbs down + db2 := s.db.WithContext(ctx).Model(&model.AgentFeedback{}).Where("status = ?", agentFeedbackStatusSubmitted) + if from != nil { + db2 = db2.Where("submitted_at >= ?", *from) + } + if to != nil { + db2 = db2.Where("submitted_at <= ?", *to) + } + db2.Where("rating = -1").Count(&stats.ThumbsDown) + + // Average dimensions via raw SQL for JSONB + type dimAvg struct { + Key string `json:"key"` + Avg float64 `json:"avg"` + } + var dimAvgs []dimAvg + rawQuery := ` + SELECT kv.key, AVG((kv.value)::numeric) AS avg + FROM agent_feedbacks, jsonb_each_text(dimensions) AS kv + WHERE status = 'submitted' AND dimensions IS NOT NULL` + args := []any{} + if from != nil { + rawQuery += " AND submitted_at >= ?" + args = append(args, *from) + } + if to != nil { + rawQuery += " AND submitted_at <= ?" + args = append(args, *to) + } + rawQuery += " GROUP BY kv.key" + s.db.WithContext(ctx).Raw(rawQuery, args...).Scan(&dimAvgs) + for _, da := range dimAvgs { + stats.AvgDimensions[da.Key] = da.Avg + } + + // Top tags + var tagCounts []TagCount + tagQuery := ` + SELECT tag AS tag, COUNT(*) AS count + FROM agent_feedbacks, jsonb_array_elements_text(tags) AS tag + WHERE status = 'submitted' AND tags IS NOT NULL` + tagArgs := []any{} + if from != nil { + tagQuery += " AND submitted_at >= ?" + tagArgs = append(tagArgs, *from) + } + if to != nil { + tagQuery += " AND submitted_at <= ?" + tagArgs = append(tagArgs, *to) + } + tagQuery += " GROUP BY tag ORDER BY count DESC LIMIT 20" + s.db.WithContext(ctx).Raw(tagQuery, tagArgs...).Scan(&tagCounts) + stats.TopTags = tagCounts + + return stats, nil +} + +// QuickSubmitFeedback creates a feedback record and immediately submits it (single operation). +// If a record already exists and is already submitted, returns the existing record (idempotent). +func (s *AgentService) QuickSubmitFeedback(ctx context.Context, fb *model.AgentFeedback) (*model.AgentFeedback, error) { + var existing model.AgentFeedback + err := s.db.WithContext(ctx). + Where("user_id = ? AND target_type = ? AND target_id = ?", fb.UserID, fb.TargetType, fb.TargetID). + First(&existing).Error + + now := time.Now() + + if err == nil { + // Already exists — if submitted, idempotent return + if existing.Status == agentFeedbackStatusSubmitted { + return &existing, nil + } + // Update rating and submit + updates := map[string]any{ + "rating": fb.Rating, + "status": agentFeedbackStatusSubmitted, + "submitted_at": &now, + "updated_at": now, + } + if len(fb.Tags) > 0 { + updates["tags"] = fb.Tags + } + if len(fb.Dimensions) > 0 { + updates["dimensions"] = fb.Dimensions + } + if fb.Comment != "" { + updates["comment"] = fb.Comment + } + if err := s.db.WithContext(ctx).Model(&existing).Updates(updates).Error; err != nil { + return nil, err + } + existing.Rating = fb.Rating + existing.Status = agentFeedbackStatusSubmitted + existing.SubmittedAt = &now + return &existing, nil + } + + if !errors.Is(err, gorm.ErrRecordNotFound) { + return nil, err + } + + // Create and immediately submit + fb.Status = agentFeedbackStatusSubmitted + fb.SubmittedAt = &now + if err := s.db.WithContext(ctx).Create(fb).Error; err != nil { + return nil, err + } + return fb, nil +} + +// EnrichFeedback updates the optional detail fields (tags, dimensions, comment) on any feedback +// record regardless of status. Never changes rating or status. Used for post-submit detail addition. +func (s *AgentService) EnrichFeedback( + ctx context.Context, + userID uint, + targetType, targetID string, + tags, dimensions datatypes.JSON, + comment string, +) (*model.AgentFeedback, error) { + var fb model.AgentFeedback + if err := s.db.WithContext(ctx). + Where("user_id = ? AND target_type = ? AND target_id = ?", userID, targetType, targetID). + First(&fb).Error; err != nil { + return nil, err + } + + now := time.Now() + updates := map[string]any{ + "enriched_at": &now, + "updated_at": now, + } + if len(tags) > 0 { + updates["tags"] = tags + } + if len(dimensions) > 0 { + updates["dimensions"] = dimensions + } + if comment != "" { + updates["comment"] = comment + } + + if err := s.db.WithContext(ctx).Model(&fb).Updates(updates).Error; err != nil { + return nil, err + } + fb.EnrichedAt = &now + return &fb, nil +} diff --git a/backend/internal/service/agent_service_test.go b/backend/internal/service/agent_service_test.go new file mode 100644 index 000000000..ac06d1a97 --- /dev/null +++ b/backend/internal/service/agent_service_test.go @@ -0,0 +1,42 @@ +package service + +import ( + "testing" + + "github.com/raids-lab/crater/internal/bizerr" +) + +func TestIsMissingAgentToolAuditColumnError(t *testing.T) { + t.Parallel() + + err := bizerr.Internal.DatabaseError.New( + `ERROR: column "execution_backend" of relation "agent_tool_calls" does not exist (SQLSTATE 42703)`, + ) + if !isMissingAgentToolAuditColumnError(err) { + t.Fatalf("expected missing audit column error to be detected") + } + + otherErr := bizerr.Internal.DatabaseError.New( + `ERROR: column "unknown_field" of relation "agent_tool_calls" does not exist (SQLSTATE 42703)`, + ) + if isMissingAgentToolAuditColumnError(otherErr) { + t.Fatalf("expected unrelated column error to be ignored") + } +} + +func TestStripUnsupportedAgentToolAuditUpdates(t *testing.T) { + t.Parallel() + + updates := map[string]any{ + "result_status": "success", + "execution_backend": "backend", + } + + compat := stripUnsupportedAgentToolAuditUpdates(updates) + if _, ok := compat["execution_backend"]; ok { + t.Fatalf("execution_backend should be stripped") + } + if compat["result_status"] != "success" { + t.Fatalf("expected non-audit fields to be preserved, got %+v", compat) + } +} diff --git a/backend/internal/service/agent_source.go b/backend/internal/service/agent_source.go new file mode 100644 index 000000000..cddeb244d --- /dev/null +++ b/backend/internal/service/agent_source.go @@ -0,0 +1,30 @@ +package service + +import "strings" + +const ( + agentSessionSourceChat = "chat" + agentSessionSourceOpsAudit = "ops_audit" + agentSessionSourceSystem = "system" + agentToolCallSourceBackend = "backend" +) + +func normalizeAgentSessionSource(source string) string { + switch strings.TrimSpace(strings.ToLower(source)) { + case agentSessionSourceOpsAudit: + return agentSessionSourceOpsAudit + case agentSessionSourceSystem: + return agentSessionSourceSystem + default: + return agentSessionSourceChat + } +} + +func normalizeAgentToolCallSource(source string) string { + switch strings.TrimSpace(strings.ToLower(source)) { + case "local": + return "local" + default: + return agentToolCallSourceBackend + } +} diff --git a/backend/internal/service/config_service.go b/backend/internal/service/config_service.go index a1cdfb15e..553c7ca7e 100644 --- a/backend/internal/service/config_service.go +++ b/backend/internal/service/config_service.go @@ -29,11 +29,7 @@ import ( // 定义掩码常量 const MaskedAPIKeyPlaceholder = "********************************************" -const ( - DefaultModelDownloadMaxConcurrent = 5 - DefaultModelDownloadWindowHours = 2 - DefaultModelDownloadMaxSuccessfulDownloads = 5 -) +const userLLMConfigKeyPrefix = "USER_LLM" // LLMConfig 结构体用于承载从数据库读取的配置 type LLMConfig struct { @@ -42,14 +38,43 @@ type LLMConfig struct { ModelName string } -// ModelDownloadLimitConfig controls model and dataset download quotas for all users. -// Only explicitly whitelisted users are exempt from these limits. -type ModelDownloadLimitConfig struct { - Enabled bool - MaxConcurrent int64 - WindowHours int64 - MaxSuccessfulDownloads int64 - WhitelistUserIDs []uint +type LLMConfigStatus struct { + Config *LLMConfig + Source string + Complete bool + HasAPIKey bool + UsingConfig bool +} + +type UserLLMConfigStatus struct { + Config *LLMConfig + Source string + Complete bool + HasAPIKey bool + UsingPersonal bool +} + +func BuildAgentLLMClientConfig(cfg *LLMConfig) map[string]any { + if cfg == nil { + return nil + } + if strings.TrimSpace(cfg.BaseURL) == "" || + strings.TrimSpace(cfg.ModelName) == "" || + strings.TrimSpace(cfg.APIKey) == "" { + return nil + } + defaultClient := map[string]any{ + "provider": "openai_compatible", + "base_url": strings.TrimSpace(cfg.BaseURL), + "model": strings.TrimSpace(cfg.ModelName), + "api_key": strings.TrimSpace(cfg.APIKey), + "temperature": 0.1, + "max_tokens": 8192, + "timeout": 120, + "streaming": true, + "stream_usage": true, + } + return map[string]any{"default": defaultClient} } // cleanBaseURL 内部辅助：清理 URL 结尾的斜杠 @@ -117,16 +142,6 @@ func (s *ConfigService) initDefaultConfigs(ctx context.Context) error { model.ConfigKeyBillingAccountIssueAmountOverrideEnabled, model.ConfigKeyBillingAccountIssuePeriodOverrideEnabled: defaultValue = "false" - case model.ConfigKeyModelDownloadLimitEnabled: - defaultValue = strconv.FormatBool(true) - case model.ConfigKeyModelDownloadMaxConcurrent: - defaultValue = strconv.Itoa(DefaultModelDownloadMaxConcurrent) - case model.ConfigKeyModelDownloadWindowHours: - defaultValue = strconv.Itoa(DefaultModelDownloadWindowHours) - case model.ConfigKeyModelDownloadMaxSuccessfulDownloads: - defaultValue = strconv.Itoa(DefaultModelDownloadMaxSuccessfulDownloads) - case model.ConfigKeyModelDownloadWhitelistUsers: - defaultValue = "[]" case model.ConfigKeyRunningSettlementIntervalMinute: defaultValue = "5" case model.ConfigKeyBillingDefaultIssueAmount: @@ -150,110 +165,78 @@ func (s *ConfigService) initDefaultConfigs(ctx context.Context) error { }) } -func (s *ConfigService) GetModelDownloadLimitConfig(ctx context.Context) (*ModelDownloadLimitConfig, error) { - configMap, err := s.getConfigs(ctx, - model.ConfigKeyModelDownloadLimitEnabled, - model.ConfigKeyModelDownloadMaxConcurrent, - model.ConfigKeyModelDownloadWindowHours, - model.ConfigKeyModelDownloadMaxSuccessfulDownloads, - model.ConfigKeyModelDownloadWhitelistUsers, - ) +// GetLLMConfig 从数据库按需读取最新配置 +func (s *ConfigService) GetLLMConfig(ctx context.Context) (*LLMConfig, error) { + status, err := s.GetLLMConfigStatus(ctx) if err != nil { return nil, err } + return status.Config, nil +} - parsePositive := func(key string, fallback int64) (int64, error) { - value := configMap[key] - if value == "" { - return fallback, nil - } - parsed, parseErr := strconv.ParseInt(value, 10, 64) - if parseErr != nil || parsed <= 0 { - return 0, bizerr.Internal.DatabaseError.New( - "invalid model download config " + key + "=" + strconv.Quote(value), - ) - } - return parsed, nil - } - - enabled := true - if value := configMap[model.ConfigKeyModelDownloadLimitEnabled]; value != "" { - enabled, err = strconv.ParseBool(value) - if err != nil { - return nil, bizerr.Internal.DatabaseError.New( - "invalid model download config " + model.ConfigKeyModelDownloadLimitEnabled + - "=" + strconv.Quote(value) + ": " + err.Error(), - ) - } - } - maxConcurrent, err := parsePositive( - model.ConfigKeyModelDownloadMaxConcurrent, DefaultModelDownloadMaxConcurrent, - ) +func (s *ConfigService) GetLLMConfigStatus(ctx context.Context) (*LLMConfigStatus, error) { + cfg, err := s.loadStoredLLMConfig(ctx) if err != nil { return nil, err } - windowHours, err := parsePositive(model.ConfigKeyModelDownloadWindowHours, DefaultModelDownloadWindowHours) - if err != nil { - return nil, err + status := &LLMConfigStatus{ + Config: cfg, + Source: "system_config", + Complete: isLLMConfigComplete(cfg), + HasAPIKey: strings.TrimSpace(cfg.APIKey) != "", + UsingConfig: isLLMConfigComplete(cfg), } - maxSuccessfulDownloads, err := parsePositive( - model.ConfigKeyModelDownloadMaxSuccessfulDownloads, DefaultModelDownloadMaxSuccessfulDownloads, - ) + return status, nil +} + +func (s *ConfigService) GetEffectiveLLMConfig(ctx context.Context, userID uint) (*LLMConfig, error) { + status, err := s.GetEffectiveLLMConfigStatus(ctx, userID) if err != nil { return nil, err } - whitelistUserIDs := make([]uint, 0) - if value := configMap[model.ConfigKeyModelDownloadWhitelistUsers]; value != "" { - if err := json.Unmarshal([]byte(value), &whitelistUserIDs); err != nil { - return nil, bizerr.Internal.DatabaseError.Wrap(err, "invalid model download whitelist config") - } - } - - return &ModelDownloadLimitConfig{ - Enabled: enabled, MaxConcurrent: maxConcurrent, - WindowHours: windowHours, MaxSuccessfulDownloads: maxSuccessfulDownloads, - WhitelistUserIDs: lo.Uniq(whitelistUserIDs), - }, nil + return status.Config, nil } -func (s *ConfigService) UpdateModelDownloadLimitConfig( - ctx context.Context, cfg ModelDownloadLimitConfig, -) error { - if cfg.MaxConcurrent <= 0 || cfg.WindowHours <= 0 || cfg.MaxSuccessfulDownloads <= 0 { - return bizerr.BadRequest.ParameterError.New("model download limits must be positive integers") - } - whitelistJSON, err := json.Marshal(lo.Uniq(cfg.WhitelistUserIDs)) +func (s *ConfigService) GetEffectiveLLMConfigStatus(ctx context.Context, userID uint) (*UserLLMConfigStatus, error) { + systemCfg, err := s.loadStoredLLMConfig(ctx) if err != nil { - return bizerr.BadRequest.ParameterError.Wrap(err, "invalid model download whitelist") + return nil, err + } + if userID == 0 { + return &UserLLMConfigStatus{ + Config: systemCfg, + Source: "system_config", + Complete: isLLMConfigComplete(systemCfg), + HasAPIKey: strings.TrimSpace(systemCfg.APIKey) != "", + UsingPersonal: false, + }, nil } - updates := map[string]string{ - model.ConfigKeyModelDownloadLimitEnabled: strconv.FormatBool(cfg.Enabled), - model.ConfigKeyModelDownloadMaxConcurrent: strconv.FormatInt(cfg.MaxConcurrent, 10), - model.ConfigKeyModelDownloadWindowHours: strconv.FormatInt(cfg.WindowHours, 10), - model.ConfigKeyModelDownloadMaxSuccessfulDownloads: strconv.FormatInt(cfg.MaxSuccessfulDownloads, 10), - model.ConfigKeyModelDownloadWhitelistUsers: string(whitelistJSON), + userCfg, hasPersonal, err := s.loadUserLLMConfig(ctx, userID) + if err != nil { + return nil, err } - return s.q.Transaction(func(tx *query.Query) error { - for key, value := range updates { - result, err := tx.SystemConfig.WithContext(ctx). - Where(tx.SystemConfig.Key.Eq(key)). - Update(tx.SystemConfig.Value, value) - if err != nil { - return err - } - if result.RowsAffected == 0 { - if err := tx.SystemConfig.WithContext(ctx).Create(&model.SystemConfig{Key: key, Value: value}); err != nil { - return err - } - } - } - return nil - }) + if !hasPersonal { + return &UserLLMConfigStatus{ + Config: systemCfg, + Source: "system_config", + Complete: isLLMConfigComplete(systemCfg), + HasAPIKey: strings.TrimSpace(systemCfg.APIKey) != "", + UsingPersonal: false, + }, nil + } + + effective := mergeLLMConfig(systemCfg, userCfg) + return &UserLLMConfigStatus{ + Config: effective, + Source: "user_config", + Complete: isLLMConfigComplete(effective), + HasAPIKey: strings.TrimSpace(effective.APIKey) != "", + UsingPersonal: true, + }, nil } -// GetLLMConfig 从数据库按需读取最新配置 -func (s *ConfigService) GetLLMConfig(ctx context.Context) (*LLMConfig, error) { +func (s *ConfigService) loadStoredLLMConfig(ctx context.Context) (*LLMConfig, error) { configMap, err := s.getConfigs(ctx, model.ConfigKeyLLMBaseURL, model.ConfigKeyLLMAPIKey, model.ConfigKeyLLMModelName) if err != nil { return nil, err @@ -273,11 +256,107 @@ func (s *ConfigService) GetLLMConfig(ctx context.Context) (*LLMConfig, error) { } } - return &LLMConfig{ + cfg := &LLMConfig{ BaseURL: configMap[model.ConfigKeyLLMBaseURL], APIKey: plainKey, ModelName: configMap[model.ConfigKeyLLMModelName], - }, nil + } + return cfg, nil +} + +func (s *ConfigService) loadUserLLMConfig(ctx context.Context, userID uint) (*LLMConfig, bool, error) { + baseURLKey, apiKeyKey, modelNameKey := userLLMConfigKeys(userID) + configMap, err := s.getConfigs(ctx, baseURLKey, apiKeyKey, modelNameKey) + if err != nil { + return nil, false, err + } + + encryptedKey := configMap[apiKeyKey] + hasPersonal := strings.TrimSpace(encryptedKey) != "" || + strings.TrimSpace(configMap[baseURLKey]) != "" || + strings.TrimSpace(configMap[modelNameKey]) != "" + if !hasPersonal { + return nil, false, nil + } + + plainKey := "" + if encryptedKey != "" { + decrypted, err := crypto.Decrypt(encryptedKey) + if err != nil { + klog.Errorf("Failed to decrypt user API Key: %v, assuming plain text or empty", err) + plainKey = encryptedKey + } else { + plainKey = decrypted + } + } + + return &LLMConfig{ + BaseURL: configMap[baseURLKey], + APIKey: plainKey, + ModelName: configMap[modelNameKey], + }, true, nil +} + +func mergeLLMConfig(systemCfg, userCfg *LLMConfig) *LLMConfig { + merged := &LLMConfig{} + if systemCfg != nil { + merged.BaseURL = systemCfg.BaseURL + merged.ModelName = systemCfg.ModelName + } + if userCfg == nil { + if systemCfg != nil { + merged.APIKey = systemCfg.APIKey + } + return merged + } + if strings.TrimSpace(userCfg.BaseURL) != "" { + merged.BaseURL = userCfg.BaseURL + } + if strings.TrimSpace(userCfg.ModelName) != "" { + merged.ModelName = userCfg.ModelName + } + merged.APIKey = userCfg.APIKey + return merged +} + +func userLLMConfigKeys(userID uint) (baseURLKey, apiKeyKey, modelNameKey string) { + prefix := fmt.Sprintf("%s:%d", userLLMConfigKeyPrefix, userID) + return prefix + ":BASE_URL", prefix + ":API_KEY", prefix + ":MODEL_NAME" +} + +func isLLMConfigComplete(cfg *LLMConfig) bool { + if cfg == nil { + return false + } + return strings.TrimSpace(cfg.BaseURL) != "" && + strings.TrimSpace(cfg.ModelName) != "" && + strings.TrimSpace(cfg.APIKey) != "" +} + +func (s *ConfigService) GetAgentLLMClientConfig(ctx context.Context) (map[string]any, error) { + status, err := s.GetLLMConfigStatus(ctx) + if err != nil { + return nil, err + } + if status == nil || !status.Complete { + return nil, bizerr.BadRequest.ParameterError.New( + "platform LLM config is incomplete: BaseURL, Model and API Key are required", + ) + } + return BuildAgentLLMClientConfig(status.Config), nil +} + +func (s *ConfigService) GetAgentLLMClientConfigForUser(ctx context.Context, userID uint) (map[string]any, error) { + status, err := s.GetEffectiveLLMConfigStatus(ctx, userID) + if err != nil { + return nil, err + } + if status == nil || !status.Complete { + return nil, bizerr.BadRequest.ParameterError.New( + "LLM config is incomplete: BaseURL, Model and API Key are required", + ) + } + return BuildAgentLLMClientConfig(status.Config), nil } // CheckLLMConnection 使用 /models 接口进行校验，并验证 ModelName 是否存在 @@ -345,6 +424,88 @@ func (s *ConfigService) CheckLLMConnection(ctx context.Context, cfg *LLMConfig) return nil } +func (s *ConfigService) UpdateUserLLMConfig( + ctx context.Context, + userID uint, + reqCfg *LLMConfig, + validate bool, +) error { + if userID == 0 { + return bizerr.BadRequest.ParameterError.New("invalid user id") + } + + oldCfg, hasOld, err := s.loadUserLLMConfig(ctx, userID) + if err != nil { + return err + } + + apiKey := strings.TrimSpace(reqCfg.APIKey) + if apiKey == MaskedAPIKeyPlaceholder { + if hasOld { + apiKey = oldCfg.APIKey + } else { + apiKey = "" + } + } + if apiKey == "" { + return bizerr.BadRequest.ParameterError.New("API Key is required for personal LLM config") + } + + effectiveCfg := &LLMConfig{ + BaseURL: strings.TrimSpace(reqCfg.BaseURL), + APIKey: apiKey, + ModelName: strings.TrimSpace(reqCfg.ModelName), + } + if effectiveCfg.BaseURL == "" || effectiveCfg.ModelName == "" { + systemCfg, loadErr := s.loadStoredLLMConfig(ctx) + if loadErr != nil { + return loadErr + } + effectiveCfg = mergeLLMConfig(systemCfg, effectiveCfg) + } + if !isLLMConfigComplete(effectiveCfg) { + return bizerr.BadRequest.ParameterError.New( + "LLM config is incomplete: BaseURL, Model and API Key are required", + ) + } + if validate { + if err := s.CheckLLMConnection(ctx, effectiveCfg); err != nil { + return bizerr.Conflict.ResourceStatusError.Wrap(err, "validation failed") + } + } + + encrypted, err := crypto.Encrypt(apiKey) + if err != nil { + return bizerr.Internal.ServiceError.Wrap(err, "failed to encrypt api key") + } + baseURLKey, apiKeyKey, modelNameKey := userLLMConfigKeys(userID) + updates := map[string]string{ + baseURLKey: effectiveCfg.BaseURL, + apiKeyKey: encrypted, + modelNameKey: effectiveCfg.ModelName, + } + return s.q.Transaction(func(tx *query.Query) error { + for key, value := range updates { + if err := tx.SystemConfig.WithContext(ctx).Save(&model.SystemConfig{Key: key, Value: value}); err != nil { + return err + } + } + return nil + }) +} + +func (s *ConfigService) ResetUserLLMConfig(ctx context.Context, userID uint) error { + if userID == 0 { + return bizerr.BadRequest.ParameterError.New("invalid user id") + } + baseURLKey, apiKeyKey, modelNameKey := userLLMConfigKeys(userID) + sc := s.q.SystemConfig + _, err := sc.WithContext(ctx). + Where(sc.Key.In(baseURLKey, apiKeyKey, modelNameKey)). + Delete() + return err +} + // SetGpuAnalysisEnabled 设置GPU分析功能的开关，并同步创建或更新定时任务的状态 func (s *ConfigService) SetGpuAnalysisEnabled(c *gin.Context, enable bool) error { var ctx = c.Request.Context() @@ -461,7 +622,12 @@ func (s *ConfigService) ResetLLMConfig(ctx context.Context) error { } // UpdateLLMConfig 更新配置 -func (s *ConfigService) UpdateLLMConfig(ctx context.Context, reqCfg *LLMConfig, validate bool) error { +func (s *ConfigService) UpdateLLMConfig( + ctx context.Context, + reqCfg *LLMConfig, + validate bool, +) error { + // 1. 处理 API Key 的更新逻辑 finalKeyToSave := "" if reqCfg.APIKey == MaskedAPIKeyPlaceholder { diff --git a/backend/internal/service/config_service_test.go b/backend/internal/service/config_service_test.go index e366bf17b..c426684b0 100644 --- a/backend/internal/service/config_service_test.go +++ b/backend/internal/service/config_service_test.go @@ -1,59 +1,12 @@ package service import ( - "reflect" "strings" "testing" - "gorm.io/driver/sqlite" - "gorm.io/gorm" - "github.com/raids-lab/crater/dao/model" - "github.com/raids-lab/crater/dao/query" ) -func TestModelDownloadLimitConfigDefaultsAndUpdate(t *testing.T) { - db, err := gorm.Open(sqlite.Open("file:model_download_limit_config?mode=memory&cache=shared"), &gorm.Config{}) - if err != nil { - t.Fatal(err) - } - if err := db.AutoMigrate(&model.SystemConfig{}, &model.PrequeueConfig{}); err != nil { - t.Fatal(err) - } - service := NewConfigService(query.Use(db)) - - cfg, err := service.GetModelDownloadLimitConfig(t.Context()) - if err != nil { - t.Fatal(err) - } - defaultConfig := ModelDownloadLimitConfig{ - Enabled: true, MaxConcurrent: 5, WindowHours: 2, MaxSuccessfulDownloads: 5, - WhitelistUserIDs: []uint{}, - } - if !reflect.DeepEqual(*cfg, defaultConfig) { - t.Fatalf("unexpected default model download limits: %+v", cfg) - } - - want := ModelDownloadLimitConfig{ - Enabled: false, MaxConcurrent: 3, WindowHours: 6, MaxSuccessfulDownloads: 11, - WhitelistUserIDs: []uint{9, 9, 11}, - } - if err := service.UpdateModelDownloadLimitConfig(t.Context(), want); err != nil { - t.Fatal(err) - } - got, err := service.GetModelDownloadLimitConfig(t.Context()) - if err != nil { - t.Fatal(err) - } - want.WhitelistUserIDs = []uint{9, 11} - if !reflect.DeepEqual(*got, want) { - t.Fatalf("updated config = %+v, want scalar values from %+v and deduplicated whitelist", *got, want) - } - if err := service.UpdateModelDownloadLimitConfig(t.Context(), ModelDownloadLimitConfig{}); err == nil { - t.Fatal("zero limits should be rejected") - } -} - func TestParsePrequeueRuntimeConfig(t *testing.T) { t.Parallel() diff --git a/backend/internal/service/gpu_analysis_service.go b/backend/internal/service/gpu_analysis_service.go index 12bf29cf0..e6655a7dd 100644 --- a/backend/internal/service/gpu_analysis_service.go +++ b/backend/internal/service/gpu_analysis_service.go @@ -36,7 +36,7 @@ import ( const ( PodAnalysisMinAge = 2 * time.Hour MetricsQueryDuration = 2 * time.Hour - LLMRequestTimeout = 420 * time.Second // 为LLM请求设置超时 + LLMRequestTimeout = 300 * time.Second // 为LLM请求设置超时 MaxQueryLength = 4000 // LLM查询的最大长度 SleepBetweenRetries = 2 * time.Second // LLM请求失败时的重试间隔 GpuAnalysisMaxAge = 7 * 24 * time.Hour // GPU分析记录的最大保留时间 @@ -181,6 +181,7 @@ func (s *GpuAnalysisService) TriggerAllJobsAnalysis(ctx context.Context) (int, e if len(gpuJobs) == 0 { klog.Info("Found running jobs, but none are using GPU resources.") + go s.cleanupStaleAnalyses(context.Background()) return 0, nil } @@ -246,6 +247,7 @@ func (s *GpuAnalysisService) startAnalysisWorker() { }() // 可以选择在任务间稍作停顿，避免对 K8s API 和 LLM 造成太大压力 + s.cleanupStaleAnalyses(context.Background()) time.Sleep(SleepBetweenRetries) } klog.Info("GPU analysis worker has stopped.") // 正常情况下不应执行到这里 @@ -659,6 +661,68 @@ func (s *GpuAnalysisService) execCommandInPod( return stdout.String(), nil } +func (s *GpuAnalysisService) cleanupStaleAnalyses(ctx context.Context) { + klog.Info("Starting cleanup of stale GPU analysis records...") + ga := s.q.GpuAnalysis + j := s.q.Job + + allAnalyses, err := ga.WithContext(ctx).Find() + if err != nil { + return + } + + if len(allAnalyses) == 0 { + return + } + + jobIDs := make(map[uint]struct{}) + for _, analysis := range allAnalyses { + jobIDs[analysis.JobID] = struct{}{} + } + + jobIDSlice := make([]uint, 0, len(jobIDs)) + for id := range jobIDs { + jobIDSlice = append(jobIDSlice, id) + } + + // 3. 一次性查询所有相关的 Job + // 使用 gen 的类型安全 Where 和 In + jobs, err := j.WithContext(ctx).Where(j.ID.In(jobIDSlice...)).Find() + if err != nil { + return + } + + // 4. 筛选出所有仍然有效的 Job ID + eligibleJobIDs := make(map[uint]struct{}) + for _, job := range jobs { + // 这里的 j 是 *model.Job 类型，可以直接调用我们之前定义的方法 + if IsEligibleForGpuAnalysis(job) { + eligibleJobIDs[job.ID] = struct{}{} + } + } + + // 5. 找出需要删除的 GpuAnalysis 记录的 ID + var analysisIDsToDelete []uint + for _, analysis := range allAnalyses { + _, isJobActive := eligibleJobIDs[analysis.JobID] + // 如果 Job ID 不在有效列表中，则说明该分析记录是过时的 + if !isJobActive && analysis.ReviewStatus != model.ReviewStatusPending { + analysisIDsToDelete = append(analysisIDsToDelete, analysis.ID) + } + } + + if len(analysisIDsToDelete) == 0 { + return + } + + // 6. 批量删除过时的记录 + // 使用 gen 的 Delete 方法 + _, err = ga.WithContext(ctx).Where(ga.ID.In(analysisIDsToDelete...)).Delete() + if err != nil { + return + } +} + func IsEligibleForGpuAnalysis(j *model.Job) bool { if j == nil { return false @@ -698,8 +762,8 @@ func (s *GpuAnalysisService) cleanOvertimeGpuAnalyses(ctx context.Context) { klog.Info("Starting cleanup of overtime GPU analysis records...") ga := s.q.GpuAnalysis cutoffTime := time.Now().Add(-GpuAnalysisMaxAge) - // 超过保留时间的记录全部软删除，不区分审核状态 - _, err := ga.WithContext(ctx).Where(ga.CreatedAt.Lt(cutoffTime)).Delete() + // 超时且已review的记录可以删除 + _, err := ga.WithContext(ctx).Where(ga.CreatedAt.Lt(cutoffTime), ga.ReviewStatus.Neq(uint8(model.ReviewStatusPending))).Delete() if err != nil { klog.Errorf("Error cleaning up overtime GPU analysis records: %v", err) return diff --git a/backend/internal/service/model_download_quota.go b/backend/internal/service/model_download_quota.go deleted file mode 100644 index da18ddd73..000000000 --- a/backend/internal/service/model_download_quota.go +++ /dev/null @@ -1,184 +0,0 @@ -// Copyright 2026 The Crater Project Team, RAIDS-Lab -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package service - -import ( - "context" - "fmt" - "slices" - "time" - - "gorm.io/gorm" - - "github.com/raids-lab/crater/dao/model" - "github.com/raids-lab/crater/dao/query" - "github.com/raids-lab/crater/internal/bizerr" -) - -// ModelDownloadQuotaService owns the complete quota reservation state machine. -// Every Kubernetes download Job receives a reservation, including Jobs created -// while quota enforcement is disabled or for an exempt user. Keeping those -// lifecycle records makes later configuration changes take effect immediately -// for already-running Jobs. -type ModelDownloadQuotaService struct { - configService *ConfigService -} - -func NewModelDownloadQuotaService(configService *ConfigService) *ModelDownloadQuotaService { - return &ModelDownloadQuotaService{configService: configService} -} - -// Reserve checks the administrator-managed limits when they apply and records -// the Job attempt in the same transaction as the download state transition. -func (s *ModelDownloadQuotaService) Reserve( - ctx context.Context, - db *gorm.DB, - userID uint, - downloadID uint, - action model.ModelDownloadSubmissionAction, -) error { - if s == nil || s.configService == nil { - return bizerr.Internal.DatabaseError.New("model download quota service is not initialized") - } - cfg, err := s.configService.GetModelDownloadLimitConfig(ctx) - if err != nil { - return bizerr.Internal.DatabaseError.Wrap(err, "get model download quota") - } - - tx := db.WithContext(ctx).Session(&gorm.Session{NewDB: true}) - exempt := !cfg.Enabled || slices.Contains(cfg.WhitelistUserIDs, userID) - if !exempt { - if err := lockModelDownloadQuota(tx, userID); err != nil { - return err - } - if err := checkModelDownloadConcurrentQuota(tx, userID, downloadID, cfg.MaxConcurrent); err != nil { - return err - } - if err := checkModelDownloadWindowQuota( - tx, userID, cfg.WindowHours, cfg.MaxSuccessfulDownloads, - ); err != nil { - return err - } - } - - submission := &model.ModelDownloadSubmission{ - UserID: userID, ModelDownloadID: downloadID, Action: action, - Status: model.ModelDownloadSubmissionReserved, - } - if err := query.Use(tx).ModelDownloadSubmission.Table("model_download_submissions"). - WithContext(ctx).Create(submission); err != nil { - return bizerr.Internal.DatabaseError.Wrap(err, "record model download submission") - } - return nil -} - -func lockModelDownloadQuota(db *gorm.DB, userID uint) error { - if db.Name() != "postgres" { - return nil - } - quotaIdentity := fmt.Sprintf("model-download-quota:%d", userID) - if err := db.Exec("SELECT pg_advisory_xact_lock(hashtextextended(?, 0))", quotaIdentity).Error; err != nil { - return bizerr.Internal.DatabaseError.Wrap(err, "lock model download quota") - } - return nil -} - -func checkModelDownloadConcurrentQuota( - db *gorm.DB, userID uint, downloadID uint, maxConcurrent int64, -) error { - activeQuery := db.Table("model_downloads AS download"). - Joins("JOIN model_download_submissions AS submission ON submission.model_download_id = download.id"). - Where("download.status IN ?", []model.ModelDownloadStatus{ - model.ModelDownloadStatusPending, model.ModelDownloadStatusDownloading, - }). - Where("submission.user_id = ? AND submission.status = ?", userID, model.ModelDownloadSubmissionReserved) - if downloadID != 0 { - activeQuery = activeQuery.Where("download.id <> ?", downloadID) - } - var activeCount int64 - if err := activeQuery.Distinct("download.id").Count(&activeCount).Error; err != nil { - return bizerr.Internal.DatabaseError.Wrap(err, "count concurrent model downloads") - } - if activeCount >= maxConcurrent { - return bizerr.RateLimit.TooManyRequests.New(fmt.Sprintf( - "This account can have at most %d pending or downloading tasks at the same time. "+ - "Wait for a task to finish or pause one before trying again.", - maxConcurrent, - )) - } - return nil -} - -func checkModelDownloadWindowQuota( - db *gorm.DB, userID uint, windowHours int64, maxSuccessfulDownloads int64, -) error { - windowStart := time.Now().Add(-time.Duration(windowHours) * time.Hour) - var windowUsageCount int64 - if err := db.Table("model_download_submissions AS submission"). - Joins("LEFT JOIN model_downloads AS download ON download.id = submission.model_download_id"). - Where("submission.user_id = ?", userID). - Where( - "(submission.status = ? AND submission.completed_at >= ?) OR "+ - "(submission.status = ? AND download.deleted_at IS NULL AND download.status IN ?)", - model.ModelDownloadSubmissionSucceeded, windowStart, - model.ModelDownloadSubmissionReserved, []model.ModelDownloadStatus{ - model.ModelDownloadStatusPending, model.ModelDownloadStatusDownloading, - }, - ). - Count(&windowUsageCount).Error; err != nil { - return bizerr.Internal.DatabaseError.Wrap(err, "count model download rolling-window usage") - } - if windowUsageCount >= maxSuccessfulDownloads { - return bizerr.RateLimit.TooManyRequests.New(fmt.Sprintf( - "This account can complete at most %d model or dataset downloads in a rolling %d-hour window. "+ - "Active downloads reserve slots; wait for one to finish, fail, or be paused before trying again.", - maxSuccessfulDownloads, windowHours, - )) - } - return nil -} - -// CompleteModelDownloadQuotaReservation converts the active reservation for a -// download into a successful rolling-window entry. The window starts when the -// download completes, rather than when the Job was submitted. -func CompleteModelDownloadQuotaReservation( - ctx context.Context, db *gorm.DB, downloadID uint, completedAt time.Time, -) error { - q := query.Use(db.WithContext(ctx).Session(&gorm.Session{NewDB: true})). - ModelDownloadSubmission.Table("model_download_submissions") - _, err := q.WithContext(ctx). - Where(q.ModelDownloadID.Eq(downloadID), q.Status.Eq(string(model.ModelDownloadSubmissionReserved))). - Updates(map[string]any{ - "status": model.ModelDownloadSubmissionSucceeded, - "completed_at": completedAt, - }) - return err -} - -// ReleaseModelDownloadQuotaReservation releases active reservations after a -// failed, paused, canceled, or otherwise unsuccessful download attempt. -func ReleaseModelDownloadQuotaReservation( - ctx context.Context, db *gorm.DB, downloadID uint, -) error { - q := query.Use(db.WithContext(ctx).Session(&gorm.Session{NewDB: true})). - ModelDownloadSubmission.Table("model_download_submissions") - _, err := q.WithContext(ctx). - Where(q.ModelDownloadID.Eq(downloadID), q.Status.Eq(string(model.ModelDownloadSubmissionReserved))). - Updates(map[string]any{ - "status": model.ModelDownloadSubmissionReleased, - "completed_at": nil, - }) - return err -} diff --git a/backend/internal/service/model_download_quota_test.go b/backend/internal/service/model_download_quota_test.go deleted file mode 100644 index 73e699a37..000000000 --- a/backend/internal/service/model_download_quota_test.go +++ /dev/null @@ -1,285 +0,0 @@ -// Copyright 2026 The Crater Project Team, RAIDS-Lab -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package service - -import ( - "errors" - "fmt" - "strings" - "testing" - "time" - - "gorm.io/driver/sqlite" - "gorm.io/gorm" - - "github.com/raids-lab/crater/dao/model" - "github.com/raids-lab/crater/dao/query" - "github.com/raids-lab/crater/internal/bizerr" -) - -func TestModelDownloadQuotaServiceEnforcesConfiguredLimits(t *testing.T) { - db, configService, quotaService := newModelDownloadQuotaTestServices(t, "configured_limits") - userID := uint(7) - limit := ModelDownloadLimitConfig{ - Enabled: true, MaxConcurrent: 2, WindowHours: 2, MaxSuccessfulDownloads: 2, - } - updateModelDownloadQuotaTestConfig(t, configService, limit) - - downloads := make([]model.ModelDownload, 0, 2) - for i, status := range []model.ModelDownloadStatus{ - model.ModelDownloadStatusPending, model.ModelDownloadStatusDownloading, - } { - download := createModelDownloadQuotaTestRecord(t, db, userID, fmt.Sprintf("active-%d", i), status) - createModelDownloadQuotaTestSubmission(t, db, userID, download.ID, model.ModelDownloadSubmissionReserved) - downloads = append(downloads, download) - } - target := createModelDownloadQuotaTestRecord(t, db, userID, "concurrent-target", model.ModelDownloadStatusPending) - err := quotaService.Reserve(t.Context(), db, userID, target.ID, model.ModelDownloadSubmissionCreate) - if !errors.Is(err, bizerr.RateLimit.Base) || !strings.Contains(err.Error(), "This account") { - t.Fatalf("configured concurrent limit should return an English rate-limit error, got %v", err) - } - - if err := db.Model(&model.ModelDownload{}).Where("creator_id = ?", userID). - Update("status", model.ModelDownloadStatusReady).Error; err != nil { - t.Fatal(err) - } - for i := range downloads { - if err := CompleteModelDownloadQuotaReservation(t.Context(), db, downloads[i].ID, time.Now()); err != nil { - t.Fatal(err) - } - } - err = quotaService.Reserve(t.Context(), db, userID, target.ID, model.ModelDownloadSubmissionRetry) - if !errors.Is(err, bizerr.RateLimit.Base) || !strings.Contains(err.Error(), "rolling 2-hour window") { - t.Fatalf("configured rolling-window limit should return an English rate-limit error, got %v", err) - } - - if err := db.Model(&model.ModelDownloadSubmission{}). - Where("user_id = ? AND status = ?", userID, model.ModelDownloadSubmissionSucceeded). - Update("completed_at", time.Now().Add(-3*time.Hour)).Error; err != nil { - t.Fatal(err) - } - if err := quotaService.Reserve( - t.Context(), db, userID, target.ID, model.ModelDownloadSubmissionResume, - ); err != nil { - t.Fatalf("successful downloads outside the configured window should not count: %v", err) - } -} - -func TestModelDownloadQuotaServiceTracksJobsAcrossConfigurationChanges(t *testing.T) { - db, configService, quotaService := newModelDownloadQuotaTestServices(t, "configuration_changes") - limit := ModelDownloadLimitConfig{ - Enabled: false, MaxConcurrent: 1, WindowHours: 2, MaxSuccessfulDownloads: 10, - } - updateModelDownloadQuotaTestConfig(t, configService, limit) - - userID := uint(7) - disabledJob := createModelDownloadQuotaTestRecord( - t, db, userID, "disabled-job", model.ModelDownloadStatusDownloading, - ) - if err := quotaService.Reserve( - t.Context(), db, userID, disabledJob.ID, model.ModelDownloadSubmissionCreate, - ); err != nil { - t.Fatalf("disabled quotas should skip checks but record the Job: %v", err) - } - assertModelDownloadQuotaTestSubmission(t, db, userID, disabledJob.ID) - - limit.Enabled = true - updateModelDownloadQuotaTestConfig(t, configService, limit) - afterEnable := createModelDownloadQuotaTestRecord(t, db, userID, "after-enable", model.ModelDownloadStatusPending) - if err := quotaService.Reserve( - t.Context(), db, userID, afterEnable.ID, model.ModelDownloadSubmissionCreate, - ); !errors.Is(err, bizerr.RateLimit.Base) { - t.Fatalf("a Job started while disabled must count after quotas are enabled, got %v", err) - } - - whitelistedUserID := uint(8) - limit.WhitelistUserIDs = []uint{whitelistedUserID} - updateModelDownloadQuotaTestConfig(t, configService, limit) - whitelistedJob := createModelDownloadQuotaTestRecord( - t, db, whitelistedUserID, "whitelisted-job", model.ModelDownloadStatusDownloading, - ) - if err := quotaService.Reserve( - t.Context(), db, whitelistedUserID, whitelistedJob.ID, model.ModelDownloadSubmissionCreate, - ); err != nil { - t.Fatalf("whitelisted users should skip checks but record the Job: %v", err) - } - assertModelDownloadQuotaTestSubmission(t, db, whitelistedUserID, whitelistedJob.ID) - - limit.WhitelistUserIDs = nil - updateModelDownloadQuotaTestConfig(t, configService, limit) - afterRemoval := createModelDownloadQuotaTestRecord( - t, db, whitelistedUserID, "after-whitelist-removal", model.ModelDownloadStatusPending, - ) - if err := quotaService.Reserve( - t.Context(), db, whitelistedUserID, afterRemoval.ID, model.ModelDownloadSubmissionCreate, - ); !errors.Is(err, bizerr.RateLimit.Base) { - t.Fatalf("a whitelisted Job must count after the user is removed from the whitelist, got %v", err) - } -} - -func TestModelDownloadQuotaServiceReserveUsesSubmissionTableFromScopedDB(t *testing.T) { - db, _, quotaService := newModelDownloadQuotaTestServices(t, "scoped_db_submission_table") - userID := uint(7) - download := createModelDownloadQuotaTestRecord( - t, db, userID, "scoped-db", model.ModelDownloadStatusPending, - ) - - err := query.Use(db).Transaction(func(tx *query.Query) error { - return quotaService.Reserve( - t.Context(), tx.ModelDownload.WithContext(t.Context()).UnderlyingDB(), - userID, download.ID, model.ModelDownloadSubmissionCreate, - ) - }) - if err != nil { - t.Fatalf("reserve with a model-download-scoped DB should use the submission table: %v", err) - } - assertModelDownloadQuotaTestSubmission(t, db, userID, download.ID) -} - -func TestModelDownloadQuotaServiceReservesAndSettlesByOperator(t *testing.T) { - db, configService, quotaService := newModelDownloadQuotaTestServices(t, "settlement_by_operator") - updateModelDownloadQuotaTestConfig(t, configService, ModelDownloadLimitConfig{ - Enabled: true, MaxConcurrent: 10, WindowHours: 2, MaxSuccessfulDownloads: 1, - }) - originalCreatorID := uint(7) - retrierID := uint(8) - - failed := createModelDownloadQuotaTestRecord(t, db, originalCreatorID, "failed", model.ModelDownloadStatusFailed) - createModelDownloadQuotaTestSubmission( - t, db, originalCreatorID, failed.ID, model.ModelDownloadSubmissionReleased, - ) - if err := quotaService.Reserve( - t.Context(), db, retrierID, failed.ID, model.ModelDownloadSubmissionRetry, - ); err != nil { - t.Fatalf("a released failed attempt should not consume the rolling window: %v", err) - } - if err := db.Model(&failed).Update("status", model.ModelDownloadStatusDownloading).Error; err != nil { - t.Fatal(err) - } - - next := createModelDownloadQuotaTestRecord(t, db, retrierID, "next", model.ModelDownloadStatusPending) - if err := quotaService.Reserve( - t.Context(), db, retrierID, next.ID, model.ModelDownloadSubmissionCreate, - ); !errors.Is(err, bizerr.RateLimit.Base) { - t.Fatalf("an active reservation should prevent rolling-window overbooking, got %v", err) - } - creatorNext := createModelDownloadQuotaTestRecord( - t, db, originalCreatorID, "creator-next", model.ModelDownloadStatusPending, - ) - if err := quotaService.Reserve( - t.Context(), db, originalCreatorID, creatorNext.ID, model.ModelDownloadSubmissionCreate, - ); err != nil { - t.Fatalf("another user's retry must not consume the original creator's quota: %v", err) - } - - if err := ReleaseModelDownloadQuotaReservation(t.Context(), db, failed.ID); err != nil { - t.Fatal(err) - } - if err := db.Model(&failed).Update("status", model.ModelDownloadStatusFailed).Error; err != nil { - t.Fatal(err) - } - if err := quotaService.Reserve( - t.Context(), db, retrierID, next.ID, model.ModelDownloadSubmissionCreate, - ); err != nil { - t.Fatalf("failure should release the reserved rolling-window slot: %v", err) - } - if err := CompleteModelDownloadQuotaReservation(t.Context(), db, next.ID, time.Now()); err != nil { - t.Fatal(err) - } - if err := db.Model(&next).Update("status", model.ModelDownloadStatusReady).Error; err != nil { - t.Fatal(err) - } - - third := createModelDownloadQuotaTestRecord(t, db, retrierID, "third", model.ModelDownloadStatusPending) - if err := quotaService.Reserve( - t.Context(), db, retrierID, third.ID, model.ModelDownloadSubmissionCreate, - ); !errors.Is(err, bizerr.RateLimit.Base) { - t.Fatalf("a successful download should consume the window from completion, got %v", err) - } -} - -func newModelDownloadQuotaTestServices( - t *testing.T, name string, -) (*gorm.DB, *ConfigService, *ModelDownloadQuotaService) { - t.Helper() - db, err := gorm.Open(sqlite.Open("file:"+name+"?mode=memory&cache=shared"), &gorm.Config{ - DisableForeignKeyConstraintWhenMigrating: true, - IgnoreRelationshipsWhenMigrating: true, - }) - if err != nil { - t.Fatal(err) - } - if err := db.AutoMigrate( - &model.SystemConfig{}, &model.PrequeueConfig{}, &model.ModelDownload{}, &model.ModelDownloadSubmission{}, - ); err != nil { - t.Fatal(err) - } - configService := NewConfigService(query.Use(db)) - return db, configService, NewModelDownloadQuotaService(configService) -} - -func updateModelDownloadQuotaTestConfig( - t *testing.T, configService *ConfigService, limit ModelDownloadLimitConfig, -) { - t.Helper() - if err := configService.UpdateModelDownloadLimitConfig(t.Context(), limit); err != nil { - t.Fatal(err) - } -} - -func createModelDownloadQuotaTestRecord( - t *testing.T, db *gorm.DB, creatorID uint, name string, status model.ModelDownloadStatus, -) model.ModelDownload { - t.Helper() - download := model.ModelDownload{ - Name: "owner/" + name, Source: model.ModelSourceModelScope, - Category: model.DownloadCategoryModel, Revision: "main", Path: "public/Models/owner/" + name, - Status: status, CreatorID: creatorID, - } - if err := db.Create(&download).Error; err != nil { - t.Fatal(err) - } - return download -} - -func createModelDownloadQuotaTestSubmission( - t *testing.T, - db *gorm.DB, - userID uint, - downloadID uint, - status model.ModelDownloadSubmissionStatus, -) { - t.Helper() - if err := db.Create(&model.ModelDownloadSubmission{ - UserID: userID, ModelDownloadID: downloadID, - Action: model.ModelDownloadSubmissionCreate, Status: status, - }).Error; err != nil { - t.Fatal(err) - } -} - -func assertModelDownloadQuotaTestSubmission(t *testing.T, db *gorm.DB, userID, downloadID uint) { - t.Helper() - var count int64 - if err := db.Model(&model.ModelDownloadSubmission{}). - Where("user_id = ? AND model_download_id = ? AND status = ?", userID, downloadID, - model.ModelDownloadSubmissionReserved). - Count(&count).Error; err != nil { - t.Fatal(err) - } - if count != 1 { - t.Fatalf("reserved submission count = %d, want 1", count) - } -} diff --git a/backend/internal/storage/file.go b/backend/internal/storage/file.go index eed53600c..b71c6b06a 100644 --- a/backend/internal/storage/file.go +++ b/backend/internal/storage/file.go @@ -17,7 +17,6 @@ import ( "github.com/raids-lab/crater/dao/model" "github.com/raids-lab/crater/dao/query" - "github.com/raids-lab/crater/internal/bizerr" "github.com/raids-lab/crater/internal/resputil" "github.com/raids-lab/crater/internal/util" "github.com/raids-lab/crater/pkg/config" @@ -531,10 +530,13 @@ func GetDatasetFiles(c *gin.Context) { return } ss := "/api/ss/dataset/" + strconv.FormatUint(uint64(datasetReq.ID), 10) - realPath, err := resolveDatasetStoragePath(URL, strings.TrimPrefix(c.Request.URL.Path, ss)) - if err != nil { - resputil.HandleError(c, bizerr.BadRequest.ParameterError.Wrap(err, err.Error())) - return + path := strings.TrimPrefix(c.Request.URL.Path, ss) + token := getFirstToken(path) + var realPath string + if token == "" { + realPath = URL + } else { + realPath = URL + "/" + strings.TrimPrefix(path, "/"+token) } // Stat the target first so we can distinguish "not found in storage" from @@ -566,18 +568,6 @@ func GetDatasetFiles(c *gin.Context) { resputil.Success(c, data) } -func resolveDatasetStoragePath(datasetURL, relativeURL string) (string, error) { - basePath := cleanURLPath(datasetURL) - relativePath := cleanURLPath(relativeURL) - if relativePath == "" { - return basePath, nil - } - if relativePath == ".." || strings.HasPrefix(relativePath, "../") { - return "", errors.New("dataset path must stay within the resource directory") - } - return urlpath.Join(basePath, relativePath), nil -} - func handleDirsList(fs webdav.FileSystem, path string) ([]Files, error) { ctx := context.Background() f, err := fs.OpenFile(ctx, path, os.O_RDONLY, 0) diff --git a/backend/internal/storage/path_test.go b/backend/internal/storage/path_test.go index 95046a37d..d083ebe07 100644 --- a/backend/internal/storage/path_test.go +++ b/backend/internal/storage/path_test.go @@ -55,40 +55,3 @@ func TestGetFirstToken(t *testing.T) { }) } } - -func TestResolveDatasetStoragePath(t *testing.T) { - t.Parallel() - - const base = "sugon-gpu-incoming/Models/Qwen/Qwen3-32B" - tests := []struct { - name string - relative string - want string - wantError bool - }{ - {name: "resource root", relative: "", want: base}, - {name: "one level", relative: "/figures", want: base + "/figures"}, - {name: "nested", relative: "/figures/examples", want: base + "/figures/examples"}, - {name: "normalize", relative: "/figures/../config", want: base + "/config"}, - {name: "reject traversal", relative: "/../../other-model", wantError: true}, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - got, err := resolveDatasetStoragePath(base, tt.relative) - if tt.wantError { - if err == nil { - t.Fatalf("resolveDatasetStoragePath(%q, %q) accepted traversal", base, tt.relative) - } - return - } - if err != nil { - t.Fatalf("resolveDatasetStoragePath(%q, %q) error = %v", base, tt.relative, err) - } - if got != tt.want { - t.Fatalf("resolveDatasetStoragePath(%q, %q) = %q, want %q", base, tt.relative, got, tt.want) - } - }) - } -} diff --git a/backend/pkg/config/config.go b/backend/pkg/config/config.go index 0b48b61da..5da8ec497 100644 --- a/backend/pkg/config/config.go +++ b/backend/pkg/config/config.go @@ -131,27 +131,10 @@ type Config struct { // Optional: If not specified, default values will be used. ModelDownload struct { // Image is the container image used for model download jobs. - // Optional: Defaults to the public Crater downloader image if not specified. + // Optional: Defaults to "crater-harbor.act.buaa.edu.cn/docker.io/python:3.11-slim" if not specified. Image string `json:"image"` - // HuggingFaceEndpoint is the Hub base URL used by download jobs. - // Optional: Defaults to the official Hugging Face Hub. - HuggingFaceEndpoint string `json:"huggingFaceEndpoint"` - // ModelScopeEndpoint is the ModelScope base URL used by download jobs. - // Optional: Defaults to the official ModelScope service. - ModelScopeEndpoint string `json:"modelScopeEndpoint"` } `json:"modelDownload"` - // ModelMetadata configures background metadata refresh endpoints and cache limits. - // Endpoints are tried in order and must be selected by each deployment administrator. - ModelMetadata struct { - HuggingFaceEndpoints []string `json:"huggingFaceEndpoints"` - ModelScopeEndpoints []string `json:"modelScopeEndpoints"` - LogoAllowedHosts []string `json:"logoAllowedHosts"` - LogicalPublicPrefix string `json:"logicalPublicPrefix"` - TimeoutSeconds int `json:"timeoutSeconds"` - MaxLogoBytes int64 `json:"maxLogoBytes"` - } `json:"modelMetadata"` - // Registry contains container registry configuration for image storage and building. // Optional: If Enable is false, registry functionality will be disabled. Registry struct { @@ -335,6 +318,18 @@ type Config struct { } `json:"normal"` } `json:"auth"` + // Agent contains configuration for the Crater Agent (Python) service. + // Optional: Defaults to http://localhost:8000 if not specified. + Agent struct { + // ServiceURL is the base URL of the Python Agent service. + // Optional: Defaults to "http://localhost:8000". + ServiceURL string `json:"serviceURL"` + // InternalToken is the shared secret used by the Python Agent service when calling + // internal tool execution endpoints exposed by the Go backend. + // Optional: Can also be supplied via CRATER_AGENT_INTERNAL_TOKEN. + InternalToken string `json:"internalToken"` + } `json:"agent"` + // SchedulerPlugins contains configuration for Kubernetes scheduler plugin integrations. // Optional: Individual plugins can be enabled/disabled independently. SchedulerPlugins struct { @@ -370,7 +365,7 @@ type Config struct { // ValidateConfig validates the configuration structure and checks for required fields // -//nolint:gocyclo // This is long but simple. +//nolint:gocyclo // Centralized config validation necessarily checks many independent sections. func (c *Config) ValidateConfig() error { var errors []string @@ -578,6 +573,8 @@ func (c *Config) logConfigWarnings() { } // PrintConfig prints the configuration in a formatted and readable way, masking sensitive information +// +//nolint:gocyclo // Config summary intentionally prints many optional sections. func (c *Config) PrintConfig() { klog.Info("=== Configuration Summary ===") @@ -612,10 +609,8 @@ func (c *Config) PrintConfig() { if c.ModelDownload.Image != "" { klog.Infof("Model Download Image: %s", c.ModelDownload.Image) } else { - klog.Info("Model Download Image: ") + klog.Info("Model Download Image: ") } - klog.Infof("Model Metadata Endpoints: HuggingFace=%d, ModelScope=%d", - len(c.HuggingFaceMetadataEndpoints()), len(c.ModelScopeMetadataEndpoints())) // Secrets klog.Infof("TLS Secrets: %s, %s", c.Secrets.TLSSecretName, c.Secrets.TLSForwardSecretName) @@ -646,6 +641,18 @@ func (c *Config) PrintConfig() { klog.Info("SMTP: Disabled") } + // Agent + if c.Agent.ServiceURL != "" { + klog.Infof("Agent Service URL: %s", c.Agent.ServiceURL) + } else { + klog.Info("Agent Service URL: ") + } + if c.Agent.InternalToken != "" { + klog.Info("Agent Internal Token: ") + } else { + klog.Info("Agent Internal Token: ") + } + // Authentication if c.Auth.LDAP.Enable { klog.Infof("LDAP Authentication: Enabled (Server: %s, BaseDN: %s)", diff --git a/backend/pkg/config/model_metadata.go b/backend/pkg/config/model_metadata.go deleted file mode 100644 index 70758d29e..000000000 --- a/backend/pkg/config/model_metadata.go +++ /dev/null @@ -1,123 +0,0 @@ -// Copyright 2026 The Crater Project Team, RAIDS-Lab -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package config - -import "strings" - -const ( - DefaultHuggingFaceEndpoint = "https://huggingface.co" - DefaultModelScopeEndpoint = "https://modelscope.cn" - DefaultMetadataTimeout = 20 - DefaultMaxLogoBytes int64 = 512 * 1024 - DefaultLogicalPublicPrefix = "public" -) - -var defaultLogoAllowedHosts = []string{ - "huggingface.co", - "cdn-avatars.huggingface.co", - "resouces.modelscope.cn", - "resources.modelscope.cn", -} - -func (c *Config) HuggingFaceDownloadEndpoint() string { - return endpointOrDefault(c.ModelDownload.HuggingFaceEndpoint, DefaultHuggingFaceEndpoint) -} - -func (c *Config) ModelScopeDownloadEndpoint() string { - return endpointOrDefault(c.ModelDownload.ModelScopeEndpoint, DefaultModelScopeEndpoint) -} - -func (c *Config) HuggingFaceMetadataEndpoints() []string { - return endpointsOrDefault(c.ModelMetadata.HuggingFaceEndpoints, c.HuggingFaceDownloadEndpoint()) -} - -func (c *Config) ModelScopeMetadataEndpoints() []string { - return endpointsOrDefault(c.ModelMetadata.ModelScopeEndpoints, c.ModelScopeDownloadEndpoint()) -} - -func (c *Config) MetadataTimeoutSeconds() int { - if c.ModelMetadata.TimeoutSeconds > 0 { - return c.ModelMetadata.TimeoutSeconds - } - return DefaultMetadataTimeout -} - -func (c *Config) MetadataMaxLogoBytes() int64 { - if c.ModelMetadata.MaxLogoBytes > 0 { - return c.ModelMetadata.MaxLogoBytes - } - return DefaultMaxLogoBytes -} - -func (c *Config) MetadataLogoAllowedHosts() []string { - hosts := normalizedStrings(c.ModelMetadata.LogoAllowedHosts) - if len(hosts) == 0 { - return append([]string(nil), defaultLogoAllowedHosts...) - } - return hosts -} - -func (c *Config) MetadataLogicalPublicPrefix() string { - prefix := strings.Trim(strings.TrimSpace(c.ModelMetadata.LogicalPublicPrefix), "/") - if prefix == "" { - return DefaultLogicalPublicPrefix - } - return prefix -} - -func endpointOrDefault(endpoint, fallback string) string { - endpoint = strings.TrimRight(strings.TrimSpace(endpoint), "/") - if endpoint == "" { - return fallback - } - return endpoint -} - -func endpointsOrDefault(endpoints []string, fallback string) []string { - result := make([]string, 0, len(endpoints)) - seen := make(map[string]struct{}, len(endpoints)) - for _, endpoint := range endpoints { - endpoint = strings.TrimRight(strings.TrimSpace(endpoint), "/") - if endpoint == "" { - continue - } - if _, exists := seen[endpoint]; exists { - continue - } - seen[endpoint] = struct{}{} - result = append(result, endpoint) - } - if len(result) == 0 { - return []string{fallback} - } - return result -} - -func normalizedStrings(values []string) []string { - result := make([]string, 0, len(values)) - seen := make(map[string]struct{}, len(values)) - for _, value := range values { - value = strings.ToLower(strings.TrimSuffix(strings.TrimSpace(value), ".")) - if value == "" { - continue - } - if _, exists := seen[value]; exists { - continue - } - seen[value] = struct{}{} - result = append(result, value) - } - return result -} diff --git a/backend/pkg/config/model_metadata_test.go b/backend/pkg/config/model_metadata_test.go deleted file mode 100644 index 6ea368ee9..000000000 --- a/backend/pkg/config/model_metadata_test.go +++ /dev/null @@ -1,34 +0,0 @@ -// Copyright 2026 The Crater Project Team, RAIDS-Lab -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package config - -import "testing" - -func TestMetadataLogoAllowedHosts(t *testing.T) { - var defaultConfig Config - defaults := defaultConfig.MetadataLogoAllowedHosts() - if len(defaults) != 4 || defaults[2] != "resouces.modelscope.cn" || defaults[3] != "resources.modelscope.cn" { - t.Fatalf("unexpected default logo hosts: %v", defaults) - } - - var customConfig Config - customConfig.ModelMetadata.LogoAllowedHosts = []string{ - " CDN.EXAMPLE. ", "cdn.example", "mirror.example", - } - custom := customConfig.MetadataLogoAllowedHosts() - if len(custom) != 2 || custom[0] != "cdn.example" || custom[1] != "mirror.example" { - t.Fatalf("unexpected normalized logo hosts: %v", custom) - } -} diff --git a/backend/pkg/packer/buildkit.go b/backend/pkg/packer/buildkit.go index e5d0a2a2f..0969befcf 100644 --- a/backend/pkg/packer/buildkit.go +++ b/backend/pkg/packer/buildkit.go @@ -169,7 +169,13 @@ func (b *imagePacker) generateBuildKitContainer( { Name: "buildkit", Image: config.GetConfig().Registry.BuildTools.Images.Buildx, - Args: setupCommands, + Command: []string{ + setupCommands[0], + setupCommands[1], + }, + Args: []string{ + setupCommands[2], + }, Env: []corev1.EnvVar{ { Name: "DOCKER_CONFIG", diff --git a/backend/pkg/packer/envd.go b/backend/pkg/packer/envd.go index 7c793ac6f..b20105a09 100644 --- a/backend/pkg/packer/envd.go +++ b/backend/pkg/packer/envd.go @@ -118,8 +118,14 @@ func (b *imagePacker) generateEnvdContainer(data *EnvdReq) []corev1.Container { { Name: "buildkit", Image: config.GetConfig().Registry.BuildTools.Images.Envd, - Args: setupCommands, - Env: envVars, + Command: []string{ + setupCommands[0], + setupCommands[1], + }, + Args: []string{ + setupCommands[2], + }, + Env: envVars, VolumeMounts: []corev1.VolumeMount{ { Name: "harborcredits", diff --git a/backend/pkg/prompts/client.go b/backend/pkg/prompts/client.go index 8a415b894..4f15a77c2 100644 --- a/backend/pkg/prompts/client.go +++ b/backend/pkg/prompts/client.go @@ -1,14 +1,23 @@ package prompts import ( + "bufio" "bytes" "context" "encoding/json" "fmt" + "io" "net/http" "strings" ) +const ( + bytesPerKiB = 1 << 10 + streamScannerInitialBufferSize = 64 * bytesPerKiB + streamScannerMaxBufferSize = bytesPerKiB * bytesPerKiB + llmErrorBodyReadLimitBytes = 4 * bytesPerKiB +) + // --- LLM 通用结构体 --- type Message struct { @@ -24,6 +33,57 @@ type LLMRequestPayload struct { Model string `json:"model"` Messages []Message `json:"messages"` ResponseFormat *ResponseFormat `json:"response_format,omitempty"` + Stream bool `json:"stream,omitempty"` + MaxTokens int `json:"max_tokens,omitempty"` +} + +func executeLLMRequest( + httpClient *http.Client, + ctx context.Context, + apiURL string, + apiKey string, + payload LLMRequestPayload, +) (string, error) { + payloadBytes, err := json.Marshal(payload) + if err != nil { + return "", fmt.Errorf("failed to marshal LLM request payload: %w", err) + } + + req, err := http.NewRequestWithContext(ctx, "POST", apiURL, bytes.NewBuffer(payloadBytes)) + if err != nil { + return "", fmt.Errorf("failed to create LLM request: %w", err) + } + req.Header.Set("Content-Type", "application/json") + if strings.TrimSpace(apiKey) != "" { + req.Header.Set("Authorization", "Bearer "+apiKey) + } + + resp, err := httpClient.Do(req) + if err != nil { + return "", fmt.Errorf("failed to execute LLM request: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + return "", fmt.Errorf("LLM API returned non-200 status code: %d", resp.StatusCode) + } + + var apiResp struct { + Choices []struct { + Message struct { + Content string `json:"content"` + } `json:"message"` + } `json:"choices"` + } + if err := json.NewDecoder(resp.Body).Decode(&apiResp); err != nil { + return "", fmt.Errorf("failed to decode LLM API response wrapper: %w", err) + } + + if len(apiResp.Choices) == 0 { + return "", fmt.Errorf("LLM response contained no choices") + } + + return strings.TrimSpace(apiResp.Choices[0].Message.Content), nil } // CallLLMAPI 是一个通用的 LLM API 调用泛型函数 @@ -44,52 +104,201 @@ func CallLLMAPI[T any]( ResponseFormat: &ResponseFormat{Type: "json_object"}, } + rawContent, err := executeLLMRequest(httpClient, ctx, apiURL, apiKey, payload) + if err != nil { + return nil, err + } + cleanedContent := CleanLLMJSONOutput(rawContent) + + var result T + if err := json.Unmarshal([]byte(cleanedContent), &result); err != nil { + return nil, fmt.Errorf("failed to unmarshal content JSON: %w (raw: %s)", err, rawContent) + } + + return &result, nil +} + +func CallLLMText( + httpClient *http.Client, + ctx context.Context, + apiURL string, + apiKey string, + modelName string, + systemPrompt, userPrompt string, +) (string, error) { + payload := LLMRequestPayload{ + Model: modelName, + Messages: []Message{ + {Role: "system", Content: systemPrompt}, + {Role: "user", Content: userPrompt}, + }, + } + return executeLLMRequest(httpClient, ctx, apiURL, apiKey, payload) +} + +func CheckLLMAvailable( + httpClient *http.Client, + ctx context.Context, + apiURL string, + apiKey string, + modelName string, +) error { + payload := LLMRequestPayload{ + Model: modelName, + Messages: []Message{ + {Role: "user", Content: "ping"}, + }, + MaxTokens: 1, + } payloadBytes, err := json.Marshal(payload) if err != nil { - return nil, fmt.Errorf("failed to marshal LLM request payload: %w", err) + return fmt.Errorf("failed to marshal LLM health check payload: %w", err) } req, err := http.NewRequestWithContext(ctx, "POST", apiURL, bytes.NewBuffer(payloadBytes)) if err != nil { - return nil, fmt.Errorf("failed to create LLM request: %w", err) + return fmt.Errorf("failed to create LLM health check request: %w", err) } req.Header.Set("Content-Type", "application/json") - req.Header.Set("Authorization", "Bearer "+apiKey) + if strings.TrimSpace(apiKey) != "" { + req.Header.Set("Authorization", "Bearer "+apiKey) + } resp, err := httpClient.Do(req) if err != nil { - return nil, fmt.Errorf("failed to execute LLM request: %w", err) + return fmt.Errorf("failed to execute LLM health check request: %w", err) } defer resp.Body.Close() if resp.StatusCode != http.StatusOK { - return nil, fmt.Errorf("LLM API returned non-200 status code: %d", resp.StatusCode) + return fmt.Errorf("LLM API returned non-200 status code: %d%s", resp.StatusCode, readLLMErrorBody(resp)) } + return nil +} - var apiResp struct { - Choices []struct { - Message struct { - Content string `json:"content"` - } `json:"message"` - } `json:"choices"` +func CallLLMTextStream( + httpClient *http.Client, + ctx context.Context, + apiURL string, + apiKey string, + modelName string, + systemPrompt, userPrompt string, + onDelta func(string) error, +) (string, error) { + req, err := newLLMStreamRequest(ctx, apiURL, apiKey, modelName, systemPrompt, userPrompt) + if err != nil { + return "", err } - if err := json.NewDecoder(resp.Body).Decode(&apiResp); err != nil { - return nil, fmt.Errorf("failed to decode LLM API response wrapper: %w", err) + + resp, err := httpClient.Do(req) + if err != nil { + return "", fmt.Errorf("failed to execute LLM stream request: %w", err) } + defer resp.Body.Close() - if len(apiResp.Choices) == 0 { - return nil, fmt.Errorf("LLM response contained no choices") + if resp.StatusCode != http.StatusOK { + return "", fmt.Errorf("LLM API returned non-200 status code: %d%s", resp.StatusCode, readLLMErrorBody(resp)) } - rawContent := apiResp.Choices[0].Message.Content - cleanedContent := CleanLLMJSONOutput(rawContent) + var full strings.Builder + scanner := bufio.NewScanner(resp.Body) + scanner.Buffer(make([]byte, 0, streamScannerInitialBufferSize), streamScannerMaxBufferSize) + for scanner.Scan() { + data, ok := parseStreamDataLine(scanner.Text()) + if !ok { + continue + } + if data == "[DONE]" { + break + } + if err := appendLLMStreamDelta(&full, data, onDelta); err != nil { + return full.String(), err + } + } + if err := scanner.Err(); err != nil { + return full.String(), fmt.Errorf("failed to read LLM stream: %w", err) + } + return strings.TrimSpace(full.String()), nil +} - var result T - if err := json.Unmarshal([]byte(cleanedContent), &result); err != nil { - return nil, fmt.Errorf("failed to unmarshal content JSON: %w (raw: %s)", err, rawContent) +func readLLMErrorBody(resp *http.Response) string { + if resp == nil || resp.Body == nil { + return "" + } + bodyBytes, err := io.ReadAll(io.LimitReader(resp.Body, llmErrorBodyReadLimitBytes)) + if err != nil { + return "" } + body := strings.TrimSpace(string(bodyBytes)) + if body == "" { + return "" + } + return ": " + body +} - return &result, nil +func newLLMStreamRequest( + ctx context.Context, + apiURL string, + apiKey string, + modelName string, + systemPrompt string, + userPrompt string, +) (*http.Request, error) { + payload := LLMRequestPayload{ + Model: modelName, + Messages: []Message{ + {Role: "system", Content: systemPrompt}, + {Role: "user", Content: userPrompt}, + }, + Stream: true, + } + payloadBytes, err := json.Marshal(payload) + if err != nil { + return nil, fmt.Errorf("failed to marshal LLM stream request payload: %w", err) + } + req, err := http.NewRequestWithContext(ctx, "POST", apiURL, bytes.NewBuffer(payloadBytes)) + if err != nil { + return nil, fmt.Errorf("failed to create LLM stream request: %w", err) + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "text/event-stream") + if strings.TrimSpace(apiKey) != "" { + req.Header.Set("Authorization", "Bearer "+apiKey) + } + return req, nil +} + +func parseStreamDataLine(rawLine string) (string, bool) { + line := strings.TrimSpace(rawLine) + if line == "" || strings.HasPrefix(line, ":") || !strings.HasPrefix(line, "data:") { + return "", false + } + return strings.TrimSpace(strings.TrimPrefix(line, "data:")), true +} + +func appendLLMStreamDelta(full *strings.Builder, data string, onDelta func(string) error) error { + var chunk struct { + Choices []struct { + Delta struct { + Content string `json:"content"` + } `json:"delta"` + } `json:"choices"` + } + if err := json.Unmarshal([]byte(data), &chunk); err != nil { + return fmt.Errorf("failed to decode LLM stream chunk: %w", err) + } + if len(chunk.Choices) == 0 { + return nil + } + delta := chunk.Choices[0].Delta.Content + if delta == "" { + return nil + } + full.WriteString(delta) + if onDelta != nil { + return onDelta(delta) + } + return nil } // CleanLLMJSONOutput 移除可能存在的 Markdown 代码块标记 diff --git a/backend/pkg/reconciler/modeldownload-reconciler.go b/backend/pkg/reconciler/modeldownload-reconciler.go index 58102d556..c720a84ed 100644 --- a/backend/pkg/reconciler/modeldownload-reconciler.go +++ b/backend/pkg/reconciler/modeldownload-reconciler.go @@ -6,7 +6,6 @@ import ( "errors" "fmt" "io" - "net/http" "regexp" "strconv" "strings" @@ -15,7 +14,6 @@ import ( "github.com/go-logr/logr" "gorm.io/datatypes" "gorm.io/gorm" - "gorm.io/gorm/clause" batchv1 "k8s.io/api/batch/v1" v1 "k8s.io/api/core/v1" k8serrors "k8s.io/apimachinery/pkg/api/errors" @@ -30,8 +28,6 @@ import ( "github.com/raids-lab/crater/dao/model" "github.com/raids-lab/crater/dao/query" - "github.com/raids-lab/crater/internal/governance/modeldataset" - "github.com/raids-lab/crater/internal/service" "github.com/raids-lab/crater/pkg/config" ) @@ -250,14 +246,9 @@ func (r *ModelDownloadReconciler) persistRepositoryMetadata( ) error { organization := strings.SplitN(download.Name, "/", 2)[0] metadata := parseRepositoryMetadata(logs) - logoURL, logoData, logoContentType, logoErr := resolveRepositoryLogo(ctx, download, metadata.LogoURL) - if logoErr != nil { - // Logo collection is best effort and must never turn a successful model - // download into a failed task. A later metadata refresh can retry it. - klog.Warningf("Failed to cache repository logo for %s: %v", download.Name, logoErr) - } updates := map[string]any{ "organization": organization, + "logo_url": metadata.LogoURL, "source_url": downloadSourceURL(download), "display_name": metadata.DisplayName, "source_description": metadata.Description, @@ -271,9 +262,7 @@ func (r *ModelDownloadReconciler) persistRepositoryMetadata( "source_login_required": metadata.LoginRequired, "source_downloads": metadata.Downloads, "source_likes": metadata.Likes, - } - if len(logoData) > 0 { - updates["logo_url"] = logoURL + "metadata_refreshed_at": time.Now(), } if metadata.UpdatedAt != "" { if updatedAt, err := time.Parse(time.RFC3339, metadata.UpdatedAt); err == nil { @@ -286,107 +275,9 @@ func (r *ModelDownloadReconciler) persistRepositoryMetadata( } } - return query.GetDB().WithContext(ctx).Transaction(func(tx *gorm.DB) error { - source := model.ModelDatasetSource{ - Provider: model.ModelDatasetProvider(download.Source), - ResourceType: model.DataType(download.Category), - RepositoryID: download.Name, - RepositoryURL: downloadSourceURL(download), - Organization: organization, - LogoURL: logoURL, - LogoData: logoData, - LogoContentType: logoContentType, - DisplayName: metadata.DisplayName, - Description: metadata.Description, - License: metadata.License, - Task: metadata.Task, - Library: metadata.Library, - ModelType: metadata.ModelType, - ParameterCount: metadata.ParameterCount, - Private: metadata.Private, - Gated: metadata.Gated, - LoginRequired: metadata.LoginRequired, - Downloads: metadata.Downloads, - Likes: metadata.Likes, - } - if value, ok := updates["source_updated_at"].(time.Time); ok { - source.SourceUpdatedAt = &value - } - if value, ok := updates["source_created_at"].(time.Time); ok { - source.SourceCreatedAt = &value - } - var persisted model.ModelDatasetSource - lookup := tx.Where( - "provider = ? AND resource_type = ? AND repository_id = ?", - source.Provider, source.ResourceType, source.RepositoryID, - ).First(&persisted) - if errors.Is(lookup.Error, gorm.ErrRecordNotFound) { - if err := tx.Create(&source).Error; err != nil { - return err - } - persisted = source - } else if lookup.Error != nil { - return lookup.Error - } else if err := tx.Model(&persisted).Updates(source).Error; err != nil { - return err - } - updates["model_dataset_source_id"] = persisted.ID - if err := tx.Model(&model.ModelDownload{}).Where("id = ?", download.ID).Updates(updates).Error; err != nil { - return err - } - download.ModelDatasetSourceID = &persisted.ID - return nil - }) -} - -func resolveRepositoryLogo( - ctx context.Context, download *model.ModelDownload, metadataLogoURL string, -) (logoURL string, logoData []byte, contentType string, err error) { - organization := strings.SplitN(download.Name, "/", 2)[0] - var cached model.ModelDatasetSource - lookup := query.GetDB().WithContext(ctx). - Where("LOWER(organization) = ? AND octet_length(logo_data) > 0", strings.ToLower(organization)). - Order("updated_at DESC"). - First(&cached) - if lookup.Error == nil { - return cached.LogoURL, cached.LogoData, cached.LogoContentType, nil - } - if !errors.Is(lookup.Error, gorm.ErrRecordNotFound) { - return "", nil, "", fmt.Errorf("load cached organization logo: %w", lookup.Error) - } - - appConfig := config.GetConfig() - httpClient := &http.Client{Timeout: time.Duration(appConfig.MetadataTimeoutSeconds()) * time.Second} - logoURL = strings.TrimSpace(metadataLogoURL) - if logoURL == "" { - switch download.Source { - case model.ModelSourceHuggingFace: - logoURL, err = modeldataset.FetchHuggingFaceAvatarURL( - ctx, httpClient, appConfig.HuggingFaceMetadataEndpoints(), organization, - ) - case model.ModelSourceModelScope: - logoURL, err = modeldataset.FetchModelScopeAvatarURL( - ctx, httpClient, downloadSourceURL(download), - ) - } - if err != nil { - return "", nil, "", err - } - } - if logoURL == "" { - return "", nil, "", nil - } - logoData, contentType, err = modeldataset.FetchSourceLogo( - ctx, - httpClient, - logoURL, - appConfig.MetadataLogoAllowedHosts(), - appConfig.MetadataMaxLogoBytes(), - ) - if err != nil { - return "", nil, "", err - } - return logoURL, logoData, contentType, nil + q := query.ModelDownload + _, err := q.WithContext(ctx).Where(q.ID.Eq(download.ID)).Updates(updates) + return err } func (r *ModelDownloadReconciler) handleJobNotFound(ctx context.Context, jobName string) (ctrl.Result, error) { @@ -455,23 +346,11 @@ func (r *ModelDownloadReconciler) getJobStatus(job *batchv1.Job) model.ModelDown func (r *ModelDownloadReconciler) updateDownloadStatus( ctx context.Context, download *model.ModelDownload, status model.ModelDownloadStatus, ) error { - db := query.ModelDownload.WithContext(ctx).UnderlyingDB().Session(&gorm.Session{NewDB: true}) - return db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { - if err := tx.Model(&model.ModelDownload{}). - Where("id = ?", download.ID). - Update("status", status).Error; err != nil { - return err - } - - switch status { - case model.ModelDownloadStatusReady: - return service.CompleteModelDownloadQuotaReservation(ctx, tx, download.ID, time.Now()) - case model.ModelDownloadStatusFailed: - return service.ReleaseModelDownloadQuotaReservation(ctx, tx, download.ID) - default: - return nil - } - }) + q := query.ModelDownload + _, err := q.WithContext(ctx). + Where(q.ID.Eq(download.ID)). + Update(q.Status, status) + return err } func (r *ModelDownloadReconciler) latestPodForJob(ctx context.Context, job *batchv1.Job) (*v1.Pod, error) { @@ -701,10 +580,6 @@ type downloadFailureRule struct { // downloadFailureRules are evaluated in order; the first matching rule wins. var downloadFailureRules = []downloadFailureRule{ - { - keywords: []string{"revision_not_found"}, - reason: "Download failed: the requested revision does not exist. Check the source branches or leave revision empty to use its default.", - }, { keywords: []string{"gated", "awaiting a review", "access to model", "you must be authenticated"}, reason: "Download failed: this repository is gated and requires authorization/login on the source site.", @@ -773,17 +648,15 @@ func formatSpeed(bytesPerSec int64) string { // source site, used as the dataset's WebURL and description fallback. func downloadSourceURL(download *model.ModelDownload) string { if download.Source == model.ModelSourceHuggingFace { - endpoint := config.GetConfig().HuggingFaceDownloadEndpoint() if download.Category == model.DownloadCategoryDataset { - return endpoint + "/datasets/" + download.Name + return "https://huggingface.co/datasets/" + download.Name } - return endpoint + "/" + download.Name + return "https://huggingface.co/" + download.Name } - endpoint := config.GetConfig().ModelScopeDownloadEndpoint() if download.Category == model.DownloadCategoryDataset { - return endpoint + "/datasets/" + download.Name + return "https://modelscope.cn/datasets/" + download.Name } - return endpoint + "/models/" + download.Name + return "https://modelscope.cn/models/" + download.Name } // datasetDescriptionForDownload builds the dataset description. It prefers the @@ -846,36 +719,9 @@ func datasetExtraForDownload( func (r *ModelDownloadReconciler) createDatasetForModel( ctx context.Context, download *model.ModelDownload, readmeDesc string, repositoryTags []string, -) error { - return query.GetDB().WithContext(ctx).Transaction(func(tx *gorm.DB) error { - txQuery := query.Use(tx) - - // More than one backend instance may reconcile the same completed Job. Lock - // every download row for this logical resource so all source/revision variants - // serialize on the same lock before checking whether the Dataset exists. - // Locking only the current download row would not protect concurrent - // HuggingFace and ModelScope downloads of the same repository name. - if _, err := txQuery.ModelDownload.WithContext(ctx). - Clauses(clause.Locking{Strength: "UPDATE"}). - Where(txQuery.ModelDownload.Name.Eq(download.Name)). - Order(txQuery.ModelDownload.ID). - Find(); err != nil { - return fmt.Errorf("failed to lock model downloads for %s: %w", download.Name, err) - } - - return r.createDatasetForModelTx(ctx, txQuery, download, readmeDesc, repositoryTags) - }) -} - -func (r *ModelDownloadReconciler) createDatasetForModelTx( - ctx context.Context, - txQuery *query.Query, - download *model.ModelDownload, - readmeDesc string, - repositoryTags []string, ) error { // Create a dataset record for the downloaded model or dataset - qDataset := txQuery.Dataset + qDataset := query.Dataset // 根据 category 确定数据类型 var dataType model.DataType @@ -890,39 +736,35 @@ func (r *ModelDownloadReconciler) createDatasetForModelTx( describe := datasetDescriptionForDownload(download, readmeDesc) sourceURL := downloadSourceURL(download) - datasetURL := r.convertToPhysicalPath(download.Path) - // The physical storage location and resource type identify the downloaded - // resource. A display name is not globally unique and may also belong to a - // user-created Dataset that must not be repurposed by this reconciler. - existingDataset, err := qDataset.WithContext(ctx). - Where(qDataset.URL.Eq(datasetURL), qDataset.Type.Eq(string(dataType))). + // Check if dataset already exists for this resource (check by name only, regardless of type) + // This prevents creating duplicate records with different types + // First check for non-deleted records + existingDataset, _ := qDataset.WithContext(ctx). + Where(qDataset.Name.Eq(download.Name)). First() - if err != nil && !errors.Is(err, gorm.ErrRecordNotFound) { - return fmt.Errorf("failed to query existing dataset: %w", err) - } if existingDataset != nil { + if existingDataset.Type != dataType { + klog.Warningf("Dataset %s exists with wrong type %s, updating to %s", download.Name, existingDataset.Type, dataType) + } extra := datasetExtraForDownload(existingDataset.Extra.Data(), download, sourceURL, repositoryTags) if _, err := qDataset.WithContext(ctx).Where(qDataset.ID.Eq(existingDataset.ID)).Updates(map[string]any{ - "describe": describe, - "extra": datatypes.NewJSONType(extra), - "size_bytes": download.SizeBytes, - "model_dataset_source_id": download.ModelDatasetSourceID, + "type": dataType, + "describe": describe, + "extra": datatypes.NewJSONType(extra), + "size_bytes": download.SizeBytes, }); err != nil { return fmt.Errorf("failed to update existing dataset metadata: %w", err) } klog.V(logVerboseLevelDebug).Infof("Dataset already exists for %s %s (dataset ID: %d)", resourceLabel, download.Name, existingDataset.ID) - return r.ensureDatasetAssociations(ctx, txQuery, existingDataset.ID, download.CreatorID) + return r.ensureDatasetAssociations(ctx, existingDataset.ID, download.CreatorID) } // Check for soft-deleted records - softDeletedDataset, err := qDataset.WithContext(ctx).Unscoped(). - Where(qDataset.URL.Eq(datasetURL), qDataset.Type.Eq(string(dataType)), qDataset.DeletedAt.IsNotNull()). + softDeletedDataset, _ := qDataset.WithContext(ctx).Unscoped(). + Where(qDataset.Name.Eq(download.Name), qDataset.DeletedAt.IsNotNull()). First() - if err != nil && !errors.Is(err, gorm.ErrRecordNotFound) { - return fmt.Errorf("failed to query soft-deleted dataset: %w", err) - } if softDeletedDataset != nil { // Restore the soft-deleted dataset @@ -936,11 +778,10 @@ func (r *ModelDownloadReconciler) createDatasetForModelTx( extra := datasetExtraForDownload(softDeletedDataset.Extra.Data(), download, sourceURL, repositoryTags) updates := map[string]any{ - "type": dataType, - "describe": describe, - "extra": datatypes.NewJSONType(extra), - "size_bytes": download.SizeBytes, - "model_dataset_source_id": download.ModelDatasetSourceID, + "type": dataType, + "describe": describe, + "extra": datatypes.NewJSONType(extra), + "size_bytes": download.SizeBytes, } if _, err := qDataset.WithContext(ctx). Where(qDataset.ID.Eq(softDeletedDataset.ID)). @@ -948,18 +789,20 @@ func (r *ModelDownloadReconciler) createDatasetForModelTx( return fmt.Errorf("failed to update restored dataset metadata: %w", err) } - return r.ensureDatasetAssociations(ctx, txQuery, softDeletedDataset.ID, download.CreatorID) + return r.ensureDatasetAssociations(ctx, softDeletedDataset.ID, download.CreatorID) } + // 将前端路径(如public/222/...)转换为物理路径(如sugon-gpu-incoming/222/...)用于存储访问 + datasetURL := r.convertToPhysicalPath(download.Path) + // Create dataset record dataset := &model.Dataset{ - Name: download.Name, - URL: datasetURL, - Describe: describe, - Type: dataType, - UserID: download.CreatorID, - SizeBytes: download.SizeBytes, - ModelDatasetSourceID: download.ModelDatasetSourceID, + Name: download.Name, + URL: datasetURL, + Describe: describe, + Type: dataType, + UserID: download.CreatorID, + SizeBytes: download.SizeBytes, Extra: datatypes.NewJSONType(model.ExtraContent{ Tags: datasetExtraForDownload(model.ExtraContent{}, download, sourceURL, repositoryTags).Tags, WebURL: &sourceURL, @@ -971,7 +814,7 @@ func (r *ModelDownloadReconciler) createDatasetForModelTx( return fmt.Errorf("failed to create dataset: %w", err) } - if err := r.ensureDatasetAssociations(ctx, txQuery, dataset.ID, download.CreatorID); err != nil { + if err := r.ensureDatasetAssociations(ctx, dataset.ID, download.CreatorID); err != nil { return err } @@ -980,9 +823,9 @@ func (r *ModelDownloadReconciler) createDatasetForModelTx( } func (r *ModelDownloadReconciler) ensureDatasetAssociations( - ctx context.Context, txQuery *query.Query, datasetID, userID uint, + ctx context.Context, datasetID, userID uint, ) error { - qUserDataset := txQuery.UserDataset + qUserDataset := query.UserDataset if _, err := qUserDataset.WithContext(ctx). Where(qUserDataset.UserID.Eq(userID), qUserDataset.DatasetID.Eq(datasetID)). First(); err != nil { @@ -996,7 +839,7 @@ func (r *ModelDownloadReconciler) ensureDatasetAssociations( } } - qAccountDataset := txQuery.AccountDataset + qAccountDataset := query.AccountDataset if _, err := qAccountDataset.WithContext(ctx). Where(qAccountDataset.AccountID.Eq(model.DefaultAccountID), qAccountDataset.DatasetID.Eq(datasetID)). First(); err != nil { diff --git a/backend/pkg/reconciler/modeldownload-reconciler_test.go b/backend/pkg/reconciler/modeldownload-reconciler_test.go index 8068ae76e..4775b8530 100644 --- a/backend/pkg/reconciler/modeldownload-reconciler_test.go +++ b/backend/pkg/reconciler/modeldownload-reconciler_test.go @@ -15,16 +15,10 @@ package reconciler import ( - "context" "strings" "testing" - "gorm.io/datatypes" - "gorm.io/driver/sqlite" - "gorm.io/gorm" - "github.com/raids-lab/crater/dao/model" - "github.com/raids-lab/crater/dao/query" ) func TestClassifyDownloadFailure(t *testing.T) { @@ -36,7 +30,6 @@ func TestClassifyDownloadFailure(t *testing.T) { {name: "gated", logs: "Access to model is restricted", want: "gated"}, {name: "authentication", logs: "HTTP 401 Unauthorized", want: "access denied"}, {name: "missing revision", logs: "revision not found (404)", want: "repository or revision not found"}, - {name: "validated missing revision", logs: "[ERROR] revision_not_found: 'main'", want: "requested revision does not exist"}, {name: "storage", logs: "write failed: no space left on device", want: "no space left"}, {name: "network", logs: "connection reset by peer", want: "network error"}, {name: "fallback", logs: "trace\ncustom downloader error\n", want: "custom downloader error"}, @@ -89,178 +82,3 @@ func TestDatasetExtraForDownloadPreservesTags(t *testing.T) { t.Fatalf("unexpected dataset extra: %#v", extra) } } - -func TestUpdateDownloadStatusSettlesQuotaReservations(t *testing.T) { - db, err := gorm.Open(sqlite.Open("file:model_download_quota_reconciler?mode=memory&cache=shared"), &gorm.Config{ - DisableForeignKeyConstraintWhenMigrating: true, - IgnoreRelationshipsWhenMigrating: true, - }) - if err != nil { - t.Fatal(err) - } - if err := db.AutoMigrate(&model.ModelDownload{}, &model.ModelDownloadSubmission{}); err != nil { - t.Fatal(err) - } - query.SetDefault(db) - reconciler := &ModelDownloadReconciler{} - - successful := model.ModelDownload{ - Name: "owner/success", Source: model.ModelSourceModelScope, - Category: model.DownloadCategoryModel, Revision: "main", Path: "public/Models/owner/success", - Status: model.ModelDownloadStatusDownloading, CreatorID: 7, - } - if err := db.Create(&successful).Error; err != nil { - t.Fatal(err) - } - if err := db.Create(&model.ModelDownloadSubmission{ - UserID: 7, ModelDownloadID: successful.ID, - Action: model.ModelDownloadSubmissionRetry, Status: model.ModelDownloadSubmissionReserved, - }).Error; err != nil { - t.Fatal(err) - } - if err := reconciler.updateDownloadStatus(t.Context(), &successful, model.ModelDownloadStatusReady); err != nil { - t.Fatal(err) - } - assertQuotaSubmissionSettlement(t, db, successful.ID, model.ModelDownloadSubmissionSucceeded, true) - - failed := model.ModelDownload{ - Name: "owner/failed", Source: model.ModelSourceModelScope, - Category: model.DownloadCategoryModel, Revision: "main", Path: "public/Models/owner/failed", - Status: model.ModelDownloadStatusDownloading, CreatorID: 7, - } - if err := db.Create(&failed).Error; err != nil { - t.Fatal(err) - } - if err := db.Create(&model.ModelDownloadSubmission{ - UserID: 7, ModelDownloadID: failed.ID, - Action: model.ModelDownloadSubmissionCreate, Status: model.ModelDownloadSubmissionReserved, - }).Error; err != nil { - t.Fatal(err) - } - if err := reconciler.updateDownloadStatus(t.Context(), &failed, model.ModelDownloadStatusFailed); err != nil { - t.Fatal(err) - } - assertQuotaSubmissionSettlement(t, db, failed.ID, model.ModelDownloadSubmissionReleased, false) -} - -func assertQuotaSubmissionSettlement( - t *testing.T, - db *gorm.DB, - downloadID uint, - wantStatus model.ModelDownloadSubmissionStatus, - wantCompletion bool, -) { - t.Helper() - var submission model.ModelDownloadSubmission - if err := db.Where("model_download_id = ?", downloadID).First(&submission).Error; err != nil { - t.Fatal(err) - } - if submission.Status != wantStatus { - t.Fatalf("submission status = %s, want %s", submission.Status, wantStatus) - } - if (submission.CompletedAt != nil) != wantCompletion { - t.Fatalf("submission completion = %v, want present=%t", submission.CompletedAt, wantCompletion) - } -} - -func TestCreateDatasetForModelDoesNotRepurposeSameNameDataset(t *testing.T) { - db := newModelDownloadDatasetTestDB(t) - - manual := model.Dataset{ - Name: "owner/resource", URL: "homes/user/private-resource", Type: model.DataTypeModel, - Describe: "user-created dataset", UserID: 9, - Extra: datatypes.NewJSONType(model.ExtraContent{Editable: true}), - } - if err := db.Create(&manual).Error; err != nil { - t.Fatal(err) - } - - download := &model.ModelDownload{ - Name: "owner/resource", Source: model.ModelSourceModelScope, - Category: model.DownloadCategoryModel, Revision: "master", - Path: "storage/Models/owner/resource", SizeBytes: 42, CreatorID: 1, - } - reconciler := &ModelDownloadReconciler{} - reconcileDownloadedDataset(t, db, reconciler, download) - reconcileDownloadedDataset(t, db, reconciler, download) - - assertDatasetCount(t, db, 2) - assertManualDatasetUnchanged(t, db, &manual) - assertDownloadedDataset(t, db, download, manual.ID) -} - -func newModelDownloadDatasetTestDB(t *testing.T) *gorm.DB { - t.Helper() - - db, err := gorm.Open(sqlite.Open("file:model_download_dataset_identity?mode=memory&cache=shared"), &gorm.Config{ - DisableForeignKeyConstraintWhenMigrating: true, - IgnoreRelationshipsWhenMigrating: true, - }) - if err != nil { - t.Fatal(err) - } - if err := db.AutoMigrate(&model.Dataset{}, &model.UserDataset{}, &model.AccountDataset{}); err != nil { - t.Fatal(err) - } - return db -} - -func reconcileDownloadedDataset( - t *testing.T, db *gorm.DB, reconciler *ModelDownloadReconciler, download *model.ModelDownload, -) { - t.Helper() - - if err := reconciler.createDatasetForModelTx( - context.Background(), query.Use(db), download, "downloaded model", []string{"llm"}, - ); err != nil { - t.Fatal(err) - } -} - -func assertDatasetCount(t *testing.T, db *gorm.DB, want int64) { - t.Helper() - - var datasetCount int64 - if err := db.Model(&model.Dataset{}).Count(&datasetCount).Error; err != nil { - t.Fatal(err) - } - if datasetCount != want { - t.Fatalf("dataset count = %d, want one manual and one downloaded Dataset", datasetCount) - } -} - -func assertManualDatasetUnchanged(t *testing.T, db *gorm.DB, manual *model.Dataset) { - t.Helper() - - var unchanged model.Dataset - if err := db.First(&unchanged, manual.ID).Error; err != nil { - t.Fatal(err) - } - if unchanged.URL != manual.URL || unchanged.Describe != manual.Describe || unchanged.Type != manual.Type { - t.Fatalf("same-name user Dataset was modified: %#v", unchanged) - } - var manualPublicLinks int64 - if err := db.Model(&model.AccountDataset{}).Where("dataset_id = ?", manual.ID).Count(&manualPublicLinks).Error; err != nil { - t.Fatal(err) - } - if manualPublicLinks != 0 { - t.Fatalf("same-name user Dataset received %d public associations", manualPublicLinks) - } -} - -func assertDownloadedDataset(t *testing.T, db *gorm.DB, download *model.ModelDownload, manualDatasetID uint) { - t.Helper() - - var downloaded model.Dataset - if err := db.Where("url = ? AND type = ?", download.Path, model.DataTypeModel).First(&downloaded).Error; err != nil { - t.Fatal(err) - } - if downloaded.ID == manualDatasetID || downloaded.Name != download.Name || downloaded.SizeBytes != download.SizeBytes { - t.Fatalf("unexpected downloaded Dataset: %#v", downloaded) - } - var publicLink model.AccountDataset - if err := db.Where("dataset_id = ? AND account_id = ?", downloaded.ID, model.DefaultAccountID). - First(&publicLink).Error; err != nil { - t.Fatalf("downloaded Dataset is not public: %v", err) - } -} diff --git a/backend/storage-server.Dockerfile b/backend/storage-server.Dockerfile index 446b4b09b..45377b4b6 100644 --- a/backend/storage-server.Dockerfile +++ b/backend/storage-server.Dockerfile @@ -15,8 +15,7 @@ RUN apk add tzdata && ln -s /usr/share/zoneinfo/Asia/Shanghai /etc/localtime ENV GIN_MODE=release COPY LICENSE NOTICE /usr/share/doc/crater/ COPY $BIN_DIR/bin-${TARGETPLATFORM//\//_}/storage-server . -COPY $BIN_DIR/bin-${TARGETPLATFORM//\//_}/model-dataset-governance . -RUN chmod +x storage-server model-dataset-governance +RUN chmod +x storage-server EXPOSE 7320 diff --git a/charts/crater/Chart.yaml b/charts/crater/Chart.yaml index 11662ec39..7d59287f4 100644 --- a/charts/crater/Chart.yaml +++ b/charts/crater/Chart.yaml @@ -15,13 +15,13 @@ type: application # This is the chart version. This version number should be incremented each time you make changes # to the chart and its templates, including the app version. # Versions are expected to follow Semantic Versioning (https://semver.org/) -version: 1.1.1 +version: 1.1.2 # This is the version number of the application being deployed. This version number should be # incremented each time you make changes to the application. Versions are not expected to # follow Semantic Versioning. They should reflect the version the application is using. # It is recommended to use it with quotes. -appVersion: "1.1.1" +appVersion: "1.1.2" # Additional metadata home: https://github.com/raids-lab/crater diff --git a/charts/crater/README.md b/charts/crater/README.md index 339af2e9f..e381ad3ba 100644 --- a/charts/crater/README.md +++ b/charts/crater/README.md @@ -1,6 +1,6 @@ # crater -![Version: 1.1.1](https://img.shields.io/badge/Version-1.1.1-informational?style=flat-square) ![Type: application](https://img.shields.io/badge/Type-application-informational?style=flat-square) ![AppVersion: 1.1.1](https://img.shields.io/badge/AppVersion-1.1.1-informational?style=flat-square) +![Version: 0.1.5](https://img.shields.io/badge/Version-0.1.5-informational?style=flat-square) ![Type: application](https://img.shields.io/badge/Type-application-informational?style=flat-square) ![AppVersion: 1.0.0](https://img.shields.io/badge/AppVersion-1.0.0-informational?style=flat-square) A comprehensive AI development platform for Kubernetes that provides GPU resource management, containerized development environments, and workflow orchestration. @@ -21,7 +21,7 @@ A comprehensive AI development platform for Kubernetes that provides GPU resourc | Key | Type | Default | Description | |-----|------|---------|-------------| | affinity | object | `{"nodeAffinity":{"preferredDuringSchedulingIgnoredDuringExecution":[{"preference":{"matchExpressions":[{"key":"nvidia.com/gpu.present","operator":"NotIn","values":["true"]}]},"weight":100}]}}` | Pod affinity configuration | -| backendConfig | object | `{"auth":{"ldap":{"alias":"","attributeMapping":{"displayName":"cn","email":"mail","username":"uid"},"enable":false,"help":"","server":{"address":"ldap://ldap.example.com:389","baseDN":"dc=example,dc=org","bindDN":"cn=admin,dc=example,dc=org","bindPassword":""},"uid":{"ldapAttribute":{"gid":"gidNumber","uid":"uidNumber"},"rid":{"offset":10000,"pgidAttribute":"primaryGroupID","sidAttribute":"objectSid"},"source":"default"}},"normal":{"allowLogin":true,"allowRegister":true},"token":{"accessTokenSecret":"example-access-token","refreshTokenSecret":"example-refresh-token"}},"enableLeaderElection":false,"modelDownload":{"huggingFaceEndpoint":"https://huggingface.co","image":"ghcr.io/raids-lab/crater-model-downloader:v1.0.0","modelScopeEndpoint":"https://modelscope.cn"},"modelMetadata":{"huggingFaceEndpoints":["https://huggingface.co"],"logicalPublicPrefix":"public","logoAllowedHosts":["huggingface.co","cdn-avatars.huggingface.co","resouces.modelscope.cn","resources.modelscope.cn"],"maxLogoBytes":524288,"modelScopeEndpoints":["https://modelscope.cn"],"timeoutSeconds":20},"port":":8088","postgres":{"TimeZone":"Asia/Shanghai","dbname":"postgres","host":"crater-postgresql.crater-system.svc.cluster.local","password":"","port":5432,"sslmode":"disable","user":"postgres"},"prometheusAPI":"http://192.168.0.1:12345","registry":{"buildTools":{"proxyConfig":{"httpProxy":null,"httpsProxy":null,"noProxy":null}},"enable":false,"harbor":{"password":"","server":"harbor.example.com","user":"admin"}},"secrets":{"imagePullSecretName":"","tlsForwardSecretName":"crater-tls-forward-secret","tlsSecretName":"crater-tls-secret"},"smtp":{"enable":false,"host":"mail.example.com","notify":"example@example.com","password":"","port":25,"user":"example"},"storage":{"prefix":{"account":"accounts","public":"public","user":"users"},"pvc":{"readOnlyMany":null,"readWriteMany":"crater-rw-storage"}}}` | Backend configuration | +| backendConfig | object | `{"auth":{"ldap":{"alias":"","attributeMapping":{"displayName":"cn","email":"mail","username":"uid"},"enable":false,"help":"","server":{"address":"ldap://ldap.example.com:389","baseDN":"dc=example,dc=org","bindDN":"cn=admin,dc=example,dc=org","bindPassword":""},"uid":{"ldapAttribute":{"gid":"gidNumber","uid":"uidNumber"},"rid":{"offset":10000,"pgidAttribute":"primaryGroupID","sidAttribute":"objectSid"},"source":"default"}},"normal":{"allowLogin":true,"allowRegister":true},"token":{"accessTokenSecret":"example-access-token","refreshTokenSecret":"example-refresh-token"}},"enableLeaderElection":false,"modelDownload":{"image":"crater-harbor.act.buaa.edu.cn/docker.io/library/python:3.11-slim"},"port":":8088","postgres":{"TimeZone":"Asia/Shanghai","dbname":"postgres","host":"crater-postgresql.crater-system.svc.cluster.local","password":"","port":5432,"sslmode":"disable","user":"postgres"},"prometheusAPI":"http://192.168.0.1:12345","registry":{"buildTools":{"proxyConfig":{"httpProxy":null,"httpsProxy":null,"noProxy":null}},"enable":false,"harbor":{"password":"","server":"harbor.example.com","user":"admin"}},"secrets":{"imagePullSecretName":"","tlsForwardSecretName":"crater-tls-forward-secret","tlsSecretName":"crater-tls-secret"},"smtp":{"enable":false,"host":"mail.example.com","notify":"example@example.com","password":"","port":25,"user":"example"},"storage":{"prefix":{"account":"accounts","public":"public","user":"users"},"pvc":{"readOnlyMany":null,"readWriteMany":"crater-rw-storage"}}}` | Backend configuration | | backendConfig.auth | object | `{"ldap":{"alias":"","attributeMapping":{"displayName":"cn","email":"mail","username":"uid"},"enable":false,"help":"","server":{"address":"ldap://ldap.example.com:389","baseDN":"dc=example,dc=org","bindDN":"cn=admin,dc=example,dc=org","bindPassword":""},"uid":{"ldapAttribute":{"gid":"gidNumber","uid":"uidNumber"},"rid":{"offset":10000,"pgidAttribute":"primaryGroupID","sidAttribute":"objectSid"},"source":"default"}},"normal":{"allowLogin":true,"allowRegister":true},"token":{"accessTokenSecret":"example-access-token","refreshTokenSecret":"example-refresh-token"}}` | Configuration for authentication methods and tokens | | backendConfig.auth.ldap | object | `{"alias":"","attributeMapping":{"displayName":"cn","email":"mail","username":"uid"},"enable":false,"help":"","server":{"address":"ldap://ldap.example.com:389","baseDN":"dc=example,dc=org","bindDN":"cn=admin,dc=example,dc=org","bindPassword":""},"uid":{"ldapAttribute":{"gid":"gidNumber","uid":"uidNumber"},"rid":{"offset":10000,"pgidAttribute":"primaryGroupID","sidAttribute":"objectSid"},"source":"default"}}` | LDAP authentication settings | | backendConfig.auth.ldap.alias | string | `""` | Short display name for this auth method in the UI (e.g., "ACT", "SJTU") The UI will append suffixes like "登录" or "统一身份认证", so keep it brief. | @@ -50,17 +50,8 @@ A comprehensive AI development platform for Kubernetes that provides GPU resourc | backendConfig.auth.token.accessTokenSecret | string | `"example-access-token"` | Secret key used to sign JWT access tokens (Required) Must be a secure, randomly generated string | | backendConfig.auth.token.refreshTokenSecret | string | `"example-refresh-token"` | Secret key used to sign JWT refresh tokens (Required) Must be a secure, randomly generated string | | backendConfig.enableLeaderElection | bool | `false` | Enable leader election for controller manager to ensure high availability Defaults to false if not specified | -| backendConfig.modelDownload | object | `{"huggingFaceEndpoint":"https://huggingface.co","image":"ghcr.io/raids-lab/crater-model-downloader:v1.0.0","modelScopeEndpoint":"https://modelscope.cn"}` | Model download functionality configurations | -| backendConfig.modelDownload.huggingFaceEndpoint | string | `"https://huggingface.co"` | Hugging Face Hub base URL used by model download jobs Set this to an administrator-approved mirror or gateway when the official Hub is unavailable | -| backendConfig.modelDownload.image | string | `"ghcr.io/raids-lab/crater-model-downloader:v1.0.0"` | Container image used for model download jobs Crater's public image pins source-client versions; deployments may mirror it internally | -| backendConfig.modelDownload.modelScopeEndpoint | string | `"https://modelscope.cn"` | ModelScope base URL used by model download jobs | -| backendConfig.modelMetadata | object | `{"huggingFaceEndpoints":["https://huggingface.co"],"logicalPublicPrefix":"public","logoAllowedHosts":["huggingface.co","cdn-avatars.huggingface.co","resouces.modelscope.cn","resources.modelscope.cn"],"maxLogoBytes":524288,"modelScopeEndpoints":["https://modelscope.cn"],"timeoutSeconds":20}` | Model and dataset source metadata refresh configuration | -| backendConfig.modelMetadata.huggingFaceEndpoints | list | `["https://huggingface.co"]` | Hugging Face-compatible endpoints tried in order for public metadata refresh | -| backendConfig.modelMetadata.logicalPublicPrefix | string | `"public"` | Logical prefix in download records before mapping to backendConfig.storage.prefix.public | -| backendConfig.modelMetadata.logoAllowedHosts | list | `["huggingface.co","cdn-avatars.huggingface.co","resouces.modelscope.cn","resources.modelscope.cn"]` | Exact HTTPS hosts allowed when caching source logos; add trusted mirror/CDN hosts as needed | -| backendConfig.modelMetadata.maxLogoBytes | int | `524288` | Maximum source logo size cached by Crater, in bytes | -| backendConfig.modelMetadata.modelScopeEndpoints | list | `["https://modelscope.cn"]` | ModelScope-compatible endpoints tried in order for public metadata refresh | -| backendConfig.modelMetadata.timeoutSeconds | int | `20` | Timeout in seconds for each source metadata request | +| backendConfig.modelDownload | object | `{"image":"crater-harbor.act.buaa.edu.cn/docker.io/library/python:3.11-slim"}` | Model download functionality configurations | +| backendConfig.modelDownload.image | string | `"crater-harbor.act.buaa.edu.cn/docker.io/library/python:3.11-slim"` | Container image used for model download jobs Can be customized for local deployments with network restrictions Default: python:3.11-slim (public Docker Hub) Internal deployment uses: crater-harbor.act.buaa.edu.cn/docker.io/library/python:3.11-slim | | backendConfig.port | string | `":8088"` | Network port that the server endpoint will listen on (Required) Must be specified for the server to start | | backendConfig.postgres | object | `{"TimeZone":"Asia/Shanghai","dbname":"postgres","host":"crater-postgresql.crater-system.svc.cluster.local","password":"","port":5432,"sslmode":"disable","user":"postgres"}` | PostgreSQL database connection configuration (Required) All fields must be specified for database connectivity | | backendConfig.postgres.TimeZone | string | `"Asia/Shanghai"` | Time zone for database connections Defaults to system time zone if not specified | @@ -83,7 +74,7 @@ A comprehensive AI development platform for Kubernetes that provides GPU resourc | backendConfig.registry.harbor.server | string | `"harbor.example.com"` | Harbor registry server URL (Required) Must be a valid Harbor instance URL | | backendConfig.registry.harbor.user | string | `"admin"` | Admin username for Harbor authentication (Required) User must have appropriate permissions in Harbor | | backendConfig.secrets | object | `{"imagePullSecretName":"","tlsForwardSecretName":"crater-tls-forward-secret","tlsSecretName":"crater-tls-secret"}` | Kubernetes secret names for various security components (Required) All secret names must correspond to existing Kubernetes secrets | -| backendConfig.secrets.imagePullSecretName | string | `""` | Name of the Kubernetes secret for pulling container images from private registries. Also applied to model download Jobs when `modelDownload.image` is private. If not specified, no image pull secret will be used. | +| backendConfig.secrets.imagePullSecretName | string | `""` | Name of the Kubernetes secret for pulling container images from private registries If not specified, no image pull secret will be used | | backendConfig.secrets.tlsForwardSecretName | string | `"crater-tls-forward-secret"` | Name of the Kubernetes secret for TLS forwarding configuration (Required) Secret must contain appropriate forwarding certificates | | backendConfig.secrets.tlsSecretName | string | `"crater-tls-secret"` | Name of the Kubernetes secret containing TLS certificates for HTTPS (Required) Secret must contain 'tls.crt' and 'tls.key' keys | | backendConfig.smtp | object | `{"enable":false,"host":"mail.example.com","notify":"example@example.com","password":"","port":25,"user":"example"}` | Configuration for email notifications via SMTP If Enable is false, email notifications will be disabled | @@ -196,29 +187,6 @@ A comprehensive AI development platform for Kubernetes that provides GPU resourc | images.nerdctl.tag | string | `"latest"` | Nerdctl image tag | | images.storage.repository | string | `"ghcr.io/raids-lab/storage-server"` | Storage server image repository | | images.storage.tag | string | `"latest"` | Storage server image tag | -| modelDatasetGovernance | object | `{"apply":false,"datasetMarkerPatterns":"","datasetsSubdirectory":"Datasets","enabled":false,"excludeDirectories":".cache,.git,.conda,node_modules,site-packages,test,tests,tmp,temp","logicalPublicPrefix":"public","maxDepth":8,"maxReadmeBytes":65536,"modelWeightPatterns":"*.safetensors,pytorch_model*.bin,model*.bin,*.gguf,tf_model.h5,flax_model.msgpack","modelsSubdirectories":[],"modelsSubdirectory":"Models","resources":{"limits":{"cpu":"500m","memory":"512Mi"},"requests":{"cpu":"50m","memory":"64Mi"}},"scanTimeout":"30m","schedule":"30 2 * * 0","storageRoot":"/crater"}` | Reconcile public model and dataset storage with database records | -| modelDatasetGovernance.apply | bool | `false` | Apply reconciliation results; keep false for the first validation run | -| modelDatasetGovernance.datasetMarkerPatterns | string | `""` | Dataset marker filename patterns; empty keeps filesystem-only dataset discovery disabled | -| modelDatasetGovernance.datasetsSubdirectory | string | `"Datasets"` | Datasets subdirectory below backendConfig.storage.prefix.public | -| modelDatasetGovernance.enabled | bool | `false` | Disabled by default because full filesystem scans must first be reviewed with dry-run | -| modelDatasetGovernance.excludeDirectories | string | `".cache,.git,.conda,node_modules,site-packages,test,tests,tmp,temp"` | Directory basenames excluded from filesystem scans | -| modelDatasetGovernance.logicalPublicPrefix | string | `"public"` | Logical public prefix stored by model download records | -| modelDatasetGovernance.maxDepth | int | `8` | Maximum directory depth scanned below each resource root | -| modelDatasetGovernance.maxReadmeBytes | int | `65536` | Maximum number of bytes imported from a local README | -| modelDatasetGovernance.modelWeightPatterns | string | `"*.safetensors,pytorch_model*.bin,model*.bin,*.gguf,tf_model.h5,flax_model.msgpack"` | Model weight filename patterns used together with config.json | -| modelDatasetGovernance.modelsSubdirectory | string | `"Models"` | Models subdirectory below backendConfig.storage.prefix.public | -| modelDatasetGovernance.modelsSubdirectories | list | `[]` | Optional model subdirectories below backendConfig.storage.prefix.public; when set, takes precedence over modelsSubdirectory | -| modelDatasetGovernance.scanTimeout | string | `"30m"` | Maximum duration of one filesystem scan | -| modelDatasetGovernance.schedule | string | `"30 2 * * 0"` | Cron schedule used after an administrator explicitly enables recurring scans | -| modelDatasetGovernance.storageRoot | string | `"/crater"` | Filesystem mount root inside the storage container | -| modelMetadataRefresh | object | `{"apply":true,"batchSize":100,"delay":"100ms","enabled":true,"extraEnv":[],"resources":{"limits":{"cpu":"500m","memory":"512Mi"},"requests":{"cpu":"50m","memory":"64Mi"}},"schedule":"15 */6 * * *","staleAfter":"168h"}` | Periodically refresh public model and dataset metadata from configured source endpoints | -| modelMetadataRefresh.apply | bool | `true` | Write refreshed metadata and cached logos to the database | -| modelMetadataRefresh.batchSize | int | `100` | Number of ready download records processed per database batch | -| modelMetadataRefresh.delay | string | `"100ms"` | Delay between source requests | -| modelMetadataRefresh.enabled | bool | `true` | Enable periodic refresh; the first run backfills metadata for existing ready downloads | -| modelMetadataRefresh.extraEnv | list | `[]` | Optional environment variables such as HTTPS_PROXY and NO_PROXY for restricted networks | -| modelMetadataRefresh.schedule | string | `"15 */6 * * *"` | Cron schedule for metadata refresh | -| modelMetadataRefresh.staleAfter | string | `"168h"` | Refresh metadata older than this duration | | namespaces | object | `{"create":true,"image":"crater-images","job":"crater-workspace"}` | Namespace configuration for crater components By default, crater components run in crater namespace, while jobs and images are in separate namespaces | | namespaces.create | bool | `true` | Whether to create namespaces along with the deployment | | namespaces.image | string | `"crater-images"` | Namespace for building images | diff --git a/charts/crater/templates/crater-agent/deployment.yaml b/charts/crater/templates/crater-agent/deployment.yaml new file mode 100644 index 000000000..4929b4088 --- /dev/null +++ b/charts/crater/templates/crater-agent/deployment.yaml @@ -0,0 +1,72 @@ +{{- if .Values.agent.enabled }} +apiVersion: apps/v1 +kind: Deployment +metadata: + name: crater-agent + namespace: {{ .Release.Namespace }} + labels: + app: crater-agent +spec: + replicas: {{ .Values.agent.replicaCount }} + selector: + matchLabels: + app: crater-agent + template: + metadata: + labels: + app: crater-agent + spec: + {{- with .Values.imagePullSecrets }} + imagePullSecrets: + {{- toYaml . | nindent 8 }} + {{- end }} + automountServiceAccountToken: false + containers: + - name: crater-agent-container + image: {{ .Values.images.agent.repository }}:{{ .Values.images.agent.tag }} + imagePullPolicy: {{ .Values.imagePullPolicy }} + env: + - name: CRATER_AGENT_CRATER_BACKEND_URL + value: http://crater-backend-svc + - name: CRATER_AGENT_CRATER_BACKEND_INTERNAL_TOKEN + valueFrom: + secretKeyRef: + name: crater-agent-secret + key: internal-token + - name: CRATER_AGENT_AGENT_INTERNAL_TOKEN + valueFrom: + secretKeyRef: + name: crater-agent-secret + key: internal-token + - name: CRATER_AGENT_DEFAULT_ORCHESTRATION_MODE + value: {{ .Values.agent.defaultOrchestrationMode | quote }} + - name: CRATER_AGENT_ALLOW_LOCAL_LLM_CONFIG_FALLBACK + value: {{ .Values.agent.allowLocalLLMConfigFallback | quote }} + - name: CRATER_AGENT_HOST + value: "0.0.0.0" + - name: CRATER_AGENT_PORT + value: {{ .Values.agent.port | quote }} + {{- with .Values.agent.env }} + {{- toYaml . | nindent 12 }} + {{- end }} + ports: + - containerPort: {{ .Values.agent.port }} + name: http + protocol: TCP + readinessProbe: + httpGet: + path: /health + port: {{ .Values.agent.port }} + initialDelaySeconds: 10 + periodSeconds: 5 + livenessProbe: + httpGet: + path: /health + port: {{ .Values.agent.port }} + initialDelaySeconds: 30 + periodSeconds: 10 + resources: + {{- toYaml .Values.agent.resources | nindent 12 }} + nodeSelector: {{ .Values.nodeSelector | toYaml | nindent 8 }} + tolerations: {{ .Values.tolerations | toYaml | nindent 8 }} +{{- end }} diff --git a/charts/crater/templates/crater-agent/secret.yaml b/charts/crater/templates/crater-agent/secret.yaml new file mode 100644 index 000000000..c0bf1dd64 --- /dev/null +++ b/charts/crater/templates/crater-agent/secret.yaml @@ -0,0 +1,12 @@ +{{- if .Values.agent.enabled }} +apiVersion: v1 +kind: Secret +metadata: + name: crater-agent-secret + namespace: {{ .Release.Namespace }} + labels: + app: crater-agent +type: Opaque +stringData: + internal-token: {{ required "agent.internalToken is required when agent.enabled is true" .Values.agent.internalToken | quote }} +{{- end }} diff --git a/charts/crater/templates/crater-agent/service.yaml b/charts/crater/templates/crater-agent/service.yaml new file mode 100644 index 000000000..fc17629e6 --- /dev/null +++ b/charts/crater/templates/crater-agent/service.yaml @@ -0,0 +1,17 @@ +{{- if .Values.agent.enabled }} +apiVersion: v1 +kind: Service +metadata: + name: crater-agent-svc + namespace: {{ .Release.Namespace }} + labels: + app: crater-agent +spec: + selector: + app: crater-agent + type: {{ .Values.agent.service.type }} + ports: + - protocol: TCP + port: {{ .Values.agent.service.port }} + targetPort: {{ .Values.agent.port }} +{{- end }} diff --git a/charts/crater/templates/crater-backend/deployment.yaml b/charts/crater/templates/crater-backend/deployment.yaml index 9d6bffdd6..cdeac683f 100644 --- a/charts/crater/templates/crater-backend/deployment.yaml +++ b/charts/crater/templates/crater-backend/deployment.yaml @@ -62,6 +62,14 @@ spec: command: - /controller - --config-file=/etc/config/config.yaml + {{- if .Values.agent.enabled }} + env: + - name: CRATER_AGENT_INTERNAL_TOKEN + valueFrom: + secretKeyRef: + name: crater-agent-secret + key: internal-token + {{- end }} volumeMounts: - name: backend-conf mountPath: /etc/config/config.yaml diff --git a/charts/crater/templates/cronjob/model-dataset-governance.yaml b/charts/crater/templates/cronjob/model-dataset-governance.yaml deleted file mode 100644 index 1321cdc74..000000000 --- a/charts/crater/templates/cronjob/model-dataset-governance.yaml +++ /dev/null @@ -1,69 +0,0 @@ -{{- if .Values.modelDatasetGovernance.enabled }} -apiVersion: batch/v1 -kind: CronJob -metadata: - name: model-dataset-governance - namespace: {{ .Values.namespaces.job }} - labels: - app: webdav-crater -spec: - schedule: {{ .Values.modelDatasetGovernance.schedule | quote }} - concurrencyPolicy: Forbid - successfulJobsHistoryLimit: 1 - failedJobsHistoryLimit: 3 - jobTemplate: - spec: - backoffLimit: 0 - template: - metadata: - labels: - app: model-dataset-governance - spec: - restartPolicy: Never - {{- with .Values.imagePullSecrets }} - imagePullSecrets: - {{- toYaml . | nindent 12 }} - {{- end }} - containers: - - name: governance - image: {{ .Values.images.storage.repository }}:{{ .Values.images.storage.tag }} - imagePullPolicy: {{ .Values.imagePullPolicy }} - command: - - /model-dataset-governance - args: - {{- if .Values.modelDatasetGovernance.apply }} - - --apply - {{- end }} - - --storage-root={{ .Values.modelDatasetGovernance.storageRoot }} - - --logical-public-prefix={{ .Values.modelDatasetGovernance.logicalPublicPrefix }} - {{- if .Values.modelDatasetGovernance.modelsSubdirectories }} - - --models-subdirectories={{ join "," .Values.modelDatasetGovernance.modelsSubdirectories }} - {{- else }} - - --models-subdirectory={{ .Values.modelDatasetGovernance.modelsSubdirectory }} - {{- end }} - - --datasets-subdirectory={{ .Values.modelDatasetGovernance.datasetsSubdirectory }} - - --max-depth={{ .Values.modelDatasetGovernance.maxDepth }} - - --exclude-directories={{ .Values.modelDatasetGovernance.excludeDirectories }} - - --model-weight-patterns={{ .Values.modelDatasetGovernance.modelWeightPatterns }} - - --dataset-marker-patterns={{ .Values.modelDatasetGovernance.datasetMarkerPatterns }} - - --max-readme-bytes={{ .Values.modelDatasetGovernance.maxReadmeBytes }} - - --scan-timeout={{ .Values.modelDatasetGovernance.scanTimeout }} - volumeMounts: - - name: ss-conf - mountPath: /etc/config/config.yaml - subPath: config.yaml - - name: storage - mountPath: {{ .Values.modelDatasetGovernance.storageRoot }} - resources: - {{- toYaml .Values.modelDatasetGovernance.resources | nindent 16 }} - volumes: - - name: storage - persistentVolumeClaim: - claimName: {{ .Values.storage.pvcName }} - - name: ss-conf - configMap: - name: ss-config - items: - - key: config.yaml - path: config.yaml -{{- end }} diff --git a/charts/crater/templates/cronjob/model-metadata-refresh.yaml b/charts/crater/templates/cronjob/model-metadata-refresh.yaml deleted file mode 100644 index ab5e4f37f..000000000 --- a/charts/crater/templates/cronjob/model-metadata-refresh.yaml +++ /dev/null @@ -1,57 +0,0 @@ -{{- if .Values.modelMetadataRefresh.enabled }} -apiVersion: batch/v1 -kind: CronJob -metadata: - name: model-metadata-refresh - namespace: {{ .Release.Namespace }} - labels: - app: crater-backend -spec: - schedule: {{ .Values.modelMetadataRefresh.schedule | quote }} - concurrencyPolicy: Forbid - successfulJobsHistoryLimit: 1 - failedJobsHistoryLimit: 3 - jobTemplate: - spec: - backoffLimit: 1 - template: - metadata: - labels: - app: model-metadata-refresh - spec: - restartPolicy: Never - {{- with .Values.imagePullSecrets }} - imagePullSecrets: - {{- toYaml . | nindent 12 }} - {{- end }} - containers: - - name: refresh - image: {{ .Values.images.backend.repository }}:{{ .Values.images.backend.tag }} - imagePullPolicy: {{ .Values.imagePullPolicy }} - command: - - /model-metadata-refresh - args: - {{- if .Values.modelMetadataRefresh.apply }} - - --apply - {{- end }} - - --batch-size={{ .Values.modelMetadataRefresh.batchSize }} - - --stale-after={{ .Values.modelMetadataRefresh.staleAfter }} - - --delay={{ .Values.modelMetadataRefresh.delay }} - {{- with .Values.modelMetadataRefresh.extraEnv }} - env: - {{- toYaml . | nindent 16 }} - {{- end }} - volumeMounts: - - name: backend-conf - mountPath: /etc/config/config.yaml - subPath: config.yaml - resources: - {{- toYaml .Values.modelMetadataRefresh.resources | nindent 16 }} - volumes: - - name: backend-conf - configMap: - name: backend-config - items: - - key: config.yaml - path: config.yaml -{{- end }} diff --git a/charts/crater/values.yaml b/charts/crater/values.yaml index f2f71a5ee..b7a0c2068 100644 --- a/charts/crater/values.yaml +++ b/charts/crater/values.yaml @@ -39,6 +39,11 @@ images: repository: "ghcr.io/raids-lab/crater-frontend" # -- Frontend service image tag tag: "latest" + agent: + # -- Agent service image repository + repository: "ghcr.io/raids-lab/crater-agent" + # -- Agent service image tag + tag: "latest" storage: # -- Storage server image repository repository: "ghcr.io/raids-lab/storage-server" @@ -80,68 +85,6 @@ images: # -- DBBackup image tag tag: "16.4" -# -- Periodically refresh public model and dataset metadata from configured source endpoints -modelMetadataRefresh: - # -- Enable periodic refresh; the first run backfills metadata for existing ready downloads - enabled: true - # -- Write refreshed metadata and cached logos to the database - apply: true - # -- Cron schedule for metadata refresh - schedule: "15 */6 * * *" - # -- Number of ready download records processed per database batch - batchSize: 100 - # -- Refresh metadata older than this duration - staleAfter: "168h" - # -- Delay between source requests - delay: "100ms" - # -- Optional environment variables such as HTTPS_PROXY and NO_PROXY for restricted networks - extraEnv: [] - resources: - requests: - cpu: 50m - memory: 64Mi - limits: - cpu: 500m - memory: 512Mi - -# -- Reconcile public model and dataset storage with database records -modelDatasetGovernance: - # -- Disabled by default because full filesystem scans must first be reviewed with dry-run - enabled: false - # -- Cron schedule used after an administrator explicitly enables recurring scans - schedule: "30 2 * * 0" - # -- Filesystem mount root inside the storage container - storageRoot: "/crater" - # -- Logical public prefix stored by model download records - logicalPublicPrefix: "public" - # -- Models subdirectory below backendConfig.storage.prefix.public - modelsSubdirectory: "Models" - # -- Optional model subdirectories below backendConfig.storage.prefix.public; when set, takes precedence over modelsSubdirectory - modelsSubdirectories: [] - # -- Datasets subdirectory below backendConfig.storage.prefix.public - datasetsSubdirectory: "Datasets" - # -- Maximum directory depth scanned below each resource root - maxDepth: 8 - # -- Directory basenames excluded from filesystem scans - excludeDirectories: ".cache,.git,.conda,node_modules,site-packages,test,tests,tmp,temp" - # -- Model weight filename patterns used together with config.json - modelWeightPatterns: "*.safetensors,pytorch_model*.bin,model*.bin,*.gguf,tf_model.h5,flax_model.msgpack" - # -- Dataset marker filename patterns; empty keeps filesystem-only dataset discovery disabled - datasetMarkerPatterns: "" - # -- Maximum number of bytes imported from a local README - maxReadmeBytes: 65536 - # -- Maximum duration of one filesystem scan - scanTimeout: "30m" - # -- Apply reconciliation results; keep false for the first validation run - apply: false - resources: - requests: - cpu: 50m - memory: 64Mi - limits: - cpu: 500m - memory: 512Mi - # -- Persistent Volume Claim configuration storage: # -- Whether to create PVC or use existing pvc. @@ -176,6 +119,35 @@ affinity: values: - "true" +# -- Crater Agent Python service configuration +agent: + # -- Whether to deploy the Crater Agent service with the platform + enabled: true + # -- Number of Agent replicas + replicaCount: 1 + # -- Agent container port + port: 8000 + # -- Shared internal token for backend <-> agent service calls; replace in production + internalToken: "" + # -- Default orchestration mode. Only "single_agent" is currently supported. + defaultOrchestrationMode: single_agent + # -- Whether Python Agent may fall back to its local llm-clients.json + allowLocalLLMConfigFallback: false + service: + # -- Agent service type + type: ClusterIP + # -- Agent service port + port: 80 + # -- Additional environment variables for the Agent container + env: [] + # -- Agent resource requests and limits + resources: + limits: + cpu: 1000m + memory: 2Gi + requests: + cpu: 100m + memory: 256Mi # -- Grafana proxy configuration # Only Grafana Pro has password-free login feature. We use Nginx proxy to support password-free login for Iframe grafanaProxy: @@ -237,40 +209,19 @@ backendConfig: # -- Model download functionality configurations modelDownload: # -- Container image used for model download jobs - # Crater's public image pins source-client versions; deployments may mirror it internally - image: "ghcr.io/raids-lab/crater-model-downloader:v1.0.0" - # -- Hugging Face Hub base URL used by model download jobs - # Set this to an administrator-approved mirror or gateway when the official Hub is unavailable - huggingFaceEndpoint: "https://huggingface.co" - # -- ModelScope base URL used by model download jobs - modelScopeEndpoint: "https://modelscope.cn" - - # -- Model and dataset source metadata refresh configuration - modelMetadata: - # -- Hugging Face-compatible endpoints tried in order for public metadata refresh - huggingFaceEndpoints: - # Example for a restricted network: "https://hf-mirror.com" - - "https://huggingface.co" - # -- ModelScope-compatible endpoints tried in order for public metadata refresh - modelScopeEndpoints: - - "https://modelscope.cn" - # -- Exact HTTPS hosts allowed when caching source logos; add trusted mirror/CDN hosts as needed - logoAllowedHosts: - - "huggingface.co" - - "cdn-avatars.huggingface.co" - - "resouces.modelscope.cn" - - "resources.modelscope.cn" - # -- Logical prefix in download records before mapping to backendConfig.storage.prefix.public - logicalPublicPrefix: "public" - # -- Timeout in seconds for each source metadata request - timeoutSeconds: 20 - # -- Maximum source logo size cached by Crater, in bytes - maxLogoBytes: 524288 + # Can be customized for local deployments with network restrictions + # Default: python:3.11-slim (public Docker Hub) + # Internal deployment uses: crater-harbor.act.buaa.edu.cn/docker.io/library/python:3.11-slim + image: "crater-harbor.act.buaa.edu.cn/docker.io/library/python:3.11-slim" # -- Endpoint URL for Prometheus API used for metrics and monitoring # If not specified, Prometheus integration will be disabled prometheusAPI: http://192.168.0.1:12345 + # -- Crater Agent service + agent: + # -- Base URL of the Python Agent service; keep cluster-internal + serviceURL: http://crater-agent-svc # -- PostgreSQL database connection configuration (Required) # All fields must be specified for database connectivity postgres: @@ -330,7 +281,6 @@ backendConfig: # Secret must contain appropriate forwarding certificates tlsForwardSecretName: crater-tls-forward-secret # -- Name of the Kubernetes secret for pulling container images from private registries - # Also applied to model download Jobs when modelDownload.image is private # If not specified, no image pull secret will be used imagePullSecretName: "" diff --git a/cli/cmd/image.go b/cli/cmd/image.go index 4fe503202..f9ef6996e 100644 --- a/cli/cmd/image.go +++ b/cli/cmd/image.go @@ -4,7 +4,6 @@ import ( "fmt" "os" "slices" - "strconv" "strings" "github.com/raids-lab/crater/cli/internal/api" @@ -15,19 +14,13 @@ import ( "github.com/spf13/cobra" ) -var ( - imageTaskWriteTypes = []string{"jupyter", "webide", "custom", "pytorch", "tensorflow"} - imageTaskFilterTypes = []string{"jupyter", "webide", "custom", "pytorch", "tensorflow", "all"} - imageVisibilityTypes = []string{"Public", "Private", "UserShare", "AccountShare"} - imageShareTypes = []string{"user", "account"} - imageBuildSources = []string{"EnvdAdvanced", "EnvdRaw"} - imageArchitectures = []string{"linux/amd64", "linux/arm64"} -) +var imageTaskTypes = []string{"jupyter", "webide", "custom", "pytorch", "tensorflow", "all"} +var imageVisibilityTypes = []string{"Public", "Private", "UserShare", "AccountShare"} var imageCmd = &cobra.Command{ Use: "image", - Short: "Manage images and image builds", - Long: "Build, upload, delete, share, and update Crater images and image build records.", + Short: "View images", + Long: "View container image lists from the active Crater platform.", RunE: func(cmd *cobra.Command, args []string) error { if len(args) > 0 { return errUnknownSubcommand(cmd, args[0]) @@ -36,238 +29,20 @@ var imageCmd = &cobra.Command{ }, } -var imageBuildCmd = imageCommandGroup("build", "Manage image builds") -var imageBuildLsCmd = &cobra.Command{Use: "ls", Short: "List image build records", Args: noArgs, RunE: runImageBuildLs} -var imageBuildPipAptCmd = &cobra.Command{Use: "pip-apt", Short: "Build an image from base image plus pip/apt packages", Args: noArgs, RunE: runImageBuildPipApt} -var imageBuildDockerfileCmd = &cobra.Command{Use: "dockerfile", Short: "Build an image from Dockerfile", Args: noArgs, RunE: runImageBuildDockerfile} -var imageBuildEnvdCmd = &cobra.Command{Use: "envd", Short: "Build an image from envd", Args: noArgs, RunE: runImageBuildEnvd} -var imageBuildRemoveCmd = &cobra.Command{Use: "remove", Short: "Cancel or remove image build records", Args: noArgs, RunE: runImageBuildRemove} -var imageBuildGetCmd = &cobra.Command{Use: "get ", Short: "Get an image build record", Args: exactArgs(1, "name"), RunE: runImageBuildGet} -var imageBuildTemplateCmd = &cobra.Command{Use: "template ", Short: "Get image build template", Args: exactArgs(1, "name"), RunE: runImageBuildTemplate} -var imageBuildPodCmd = &cobra.Command{Use: "pod ", Short: "Get image build pod", Args: exactArgs(1, "id"), RunE: runImageBuildPod} - -var imageLsCmd = &cobra.Command{Use: "ls", Short: "List images", Args: noArgs, RunE: runImageLs} -var imageUploadCmd = &cobra.Command{Use: "upload", Short: "Upload/register an existing image link", Args: noArgs, RunE: runImageUpload} -var imageDeleteCmd = &cobra.Command{Use: "delete ", Short: "Delete an image", Args: exactArgs(1, "id"), RunE: runImageDelete} -var imageDeleteManyCmd = &cobra.Command{Use: "delete-many", Short: "Delete multiple images", Args: noArgs, RunE: runImageDeleteMany} -var imageDescriptionCmd = &cobra.Command{Use: "description ", Short: "Update image description", Args: exactArgs(1, "id"), RunE: runImageDescription} -var imageTypeCmd = &cobra.Command{Use: "type ", Short: "Update image task type", Args: exactArgs(1, "id"), RunE: runImageType} -var imageTagsCmd = &cobra.Command{Use: "tags ", Short: "Update image tags", Args: exactArgs(1, "id"), RunE: runImageTags} -var imageArchCmd = &cobra.Command{Use: "arch ", Short: "Update image architectures", Args: exactArgs(1, "id"), RunE: runImageArch} -var imageValidCmd = &cobra.Command{Use: "valid", Short: "Validate image links", Args: noArgs, RunE: runImageValid} - -var imageShareCmd = imageCommandGroup("share", "Manage image sharing") -var imageShareLsCmd = &cobra.Command{Use: "ls ", Short: "List image grants", Args: exactArgs(1, "image-id"), RunE: runImageShareLs} -var imageShareUsersCmd = &cobra.Command{Use: "users ", Short: "List users not granted an image", Args: exactArgs(1, "image-id"), RunE: runImageShareUsers} -var imageShareAccountsCmd = &cobra.Command{Use: "accounts ", Short: "List accounts not granted an image", Args: exactArgs(1, "image-id"), RunE: runImageShareAccounts} -var imageShareAddCmd = &cobra.Command{Use: "add ", Short: "Share an image with users or accounts", Args: exactArgs(1, "image-id"), RunE: runImageShareAdd} -var imageShareRemoveCmd = &cobra.Command{Use: "remove ", Short: "Cancel image sharing", Args: exactArgs(1, "image-id"), RunE: runImageShareRemove} - -var imageCudaCmd = imageCommandGroup("cuda", "View CUDA base images") -var imageCudaLsCmd = &cobra.Command{Use: "ls", Short: "List CUDA base images", Args: noArgs, RunE: runImageCudaLs} - -var imageHarborCmd = imageCommandGroup("harbor", "View Harbor information") -var imageHarborInfoCmd = &cobra.Command{Use: "info", Short: "Get Harbor address", Args: noArgs, RunE: runImageHarborInfo} -var imageHarborCredentialCmd = &cobra.Command{Use: "credential", Short: "Create and show Harbor project credentials", Args: noArgs, RunE: runImageHarborCredential} - -var imageQuotaCmd = imageCommandGroup("quota", "View or update Harbor project quota") -var imageQuotaGetCmd = &cobra.Command{Use: "get", Short: "Get Harbor project quota", Args: noArgs, RunE: runImageQuotaGet} -var imageQuotaSetCmd = &cobra.Command{Use: "set", Short: "Update Harbor project quota", Args: noArgs, RunE: runImageQuotaSet} - -var adminImageCmd = imageCommandGroup("image", "Manage admin image resources") -var adminImageBuildLsCmd = &cobra.Command{Use: "build-ls", Short: "List all image build records", Args: noArgs, RunE: runAdminImageBuildLs} -var adminImageBuildRemoveCmd = &cobra.Command{Use: "build-remove", Short: "Cancel or remove image build records", Args: noArgs, RunE: runAdminImageBuildRemove} -var adminImageLsCmd = &cobra.Command{Use: "ls", Short: "List all images", Args: noArgs, RunE: runAdminImageLs} -var adminImageDeleteManyCmd = &cobra.Command{Use: "delete-many", Short: "Delete multiple images", Args: noArgs, RunE: runAdminImageDeleteMany} -var adminImageDescriptionCmd = &cobra.Command{Use: "description ", Short: "Update image description", Args: exactArgs(1, "id"), RunE: runAdminImageDescription} -var adminImageTypeCmd = &cobra.Command{Use: "type ", Short: "Update image task type", Args: exactArgs(1, "id"), RunE: runAdminImageType} -var adminImageTagsCmd = &cobra.Command{Use: "tags ", Short: "Update image tags", Args: exactArgs(1, "id"), RunE: runAdminImageTags} -var adminImageArchCmd = &cobra.Command{Use: "arch ", Short: "Update image architectures", Args: exactArgs(1, "id"), RunE: runAdminImageArch} -var adminImagePublicCmd = &cobra.Command{Use: "public ", Short: "Toggle public visibility", Args: exactArgs(1, "id"), RunE: runAdminImagePublic} -var adminImageCudaCmd = imageCommandGroup("cuda", "Manage CUDA base images") -var adminImageCudaAddCmd = &cobra.Command{Use: "add", Short: "Add a CUDA base image", Args: noArgs, RunE: runAdminImageCudaAdd} -var adminImageCudaDeleteCmd = &cobra.Command{Use: "delete ", Short: "Delete a CUDA base image", Args: exactArgs(1, "id"), RunE: runAdminImageCudaDelete} - -func imageCommandGroup(use, short string) *cobra.Command { - return &cobra.Command{ - Use: use, - Short: short, - RunE: func(cmd *cobra.Command, args []string) error { - if len(args) > 0 { - return errUnknownSubcommand(cmd, args[0]) - } - return cmd.Help() - }, - } -} - -func runImageBuildLs(cmd *cobra.Command, _ []string) error { - client, err := activeAPIClient() - if err != nil { - return err - } - resp, err := client.ListKaniko(false) - if err != nil { - return cliErrFromAPI(err) - } - if outputJSON { - return output.WriteSuccessJSON(os.Stdout, output.SuccessEnvelope(map[string]interface{}{"builds": resp.KanikoList})) - } - printKanikoTable(resp.KanikoList) - return nil -} - -func runAdminImageBuildLs(cmd *cobra.Command, _ []string) error { - return runImageBuildList(cmd, true) -} - -func runImageBuildList(_ *cobra.Command, admin bool) error { - client, err := activeAPIClient() - if err != nil { - return err - } - resp, err := client.ListKaniko(admin) - if err != nil { - return cliErrFromAPI(err) - } - if outputJSON { - return output.WriteSuccessJSON(os.Stdout, output.SuccessEnvelope(map[string]interface{}{"builds": resp.KanikoList})) - } - printKanikoTable(resp.KanikoList) - return nil -} - -func runImageBuildPipApt(cmd *cobra.Command, _ []string) error { - req, err := collectPipAptBuild(cmd) - if err != nil { - return err - } - client, err := activeAPIClient() - if err != nil { - return err - } - msg, err := client.CreatePipApt(req) - return writeImageMessage(msg, err) -} - -func runImageBuildDockerfile(cmd *cobra.Command, _ []string) error { - req, err := collectDockerfileBuild(cmd) - if err != nil { - return err - } - client, err := activeAPIClient() - if err != nil { - return err - } - msg, err := client.CreateDockerfile(req) - return writeImageMessage(msg, err) -} - -func runImageBuildEnvd(cmd *cobra.Command, _ []string) error { - req, err := collectEnvdBuild(cmd) - if err != nil { - return err - } - client, err := activeAPIClient() - if err != nil { - return err - } - msg, err := client.CreateEnvd(req) - return writeImageMessage(msg, err) -} - -func runImageBuildRemove(cmd *cobra.Command, _ []string) error { - ids, err := idsFlag(cmd) - if err != nil { - return err - } - client, err := activeAPIClient() - if err != nil { - return err - } - msg, err := client.RemoveKaniko(ids, false) - return writeImageMessage(msg, err) -} - -func runAdminImageBuildRemove(cmd *cobra.Command, _ []string) error { - ids, err := idsFlag(cmd) - if err != nil { - return err - } - client, err := activeAPIClient() - if err != nil { - return err - } - msg, err := client.RemoveKaniko(ids, true) - return writeImageMessage(msg, err) -} - -func runImageBuildGet(_ *cobra.Command, args []string) error { - client, err := activeAPIClient() - if err != nil { - return err - } - resp, err := client.GetKanikoByName(args[0]) - if err != nil { - return cliErrFromAPI(err) - } - if outputJSON { - return output.WriteSuccessJSON(os.Stdout, output.SuccessEnvelope(map[string]interface{}{"build": resp})) - } - printRawObject(map[string]interface{}{ - "ID": resp.ID, - "imageLink": resp.ImageLink, - "status": resp.Status, - "imagepackName": resp.ImagePackName, - "podName": resp.PodName, - "namespace": resp.PodNameSpace, - "nodeName": resp.NodeName, - }) - return nil -} - -func runImageBuildTemplate(_ *cobra.Command, args []string) error { - client, err := activeAPIClient() - if err != nil { - return err - } - template, err := client.GetKanikoTemplateByName(args[0]) - if err != nil { - return cliErrFromAPI(err) - } - if outputJSON { - return output.WriteSuccessJSON(os.Stdout, output.SuccessEnvelope(map[string]interface{}{"template": template})) - } - fmt.Println(template) - return nil -} - -func runImageBuildPod(_ *cobra.Command, args []string) error { - id, err := requiredUintArg(args, "image_label_id", "id") - if err != nil { - return err - } - client, err := activeAPIClient() - if err != nil { - return err - } - resp, err := client.GetKanikoPod(id) - if err != nil { - return cliErrFromAPI(err) - } - if outputJSON { - return output.WriteSuccessJSON(os.Stdout, output.SuccessEnvelope(map[string]interface{}{"pod": resp})) - } - printRawObject(map[string]interface{}{"name": resp.Name, "namespace": resp.Namespace, "nodeName": resp.NodeName}) - return nil +var imageLsCmd = &cobra.Command{ + Use: "ls", + Short: "List images", + Args: noArgs, + RunE: runImageLs, } func runImageLs(cmd *cobra.Command, _ []string) error { + available, _ := cmd.Flags().GetBool("available") taskType, _ := cmd.Flags().GetString("type") visibility, _ := cmd.Flags().GetString("visibility") taskType = strings.TrimSpace(taskType) visibility = strings.TrimSpace(visibility) - if taskType != "" && !slices.Contains(imageTaskFilterTypes, taskType) { + if taskType != "" && !slices.Contains(imageTaskTypes, taskType) { return errUsageFromIssues([]usageIssue{{ Code: errorcodes.ErrInvalidFlagValue, Message: i18n.T("err_invalid_image_type", taskType), @@ -285,43 +60,18 @@ func runImageLs(cmd *cobra.Command, _ []string) error { if err != nil { return err } - available, _ := cmd.Flags().GetBool("available") - var images []api.ImageInfo - if available { - images, err = client.ListAvailableImages() - } else { - var resp *api.ListImageResponse - resp, err = client.ListImageRecords(false) - if resp != nil { - images = resp.ImageList - } - } + images, err := client.ListImages(available) if err != nil { return cliErrFromAPI(err) } - if taskType != "" && taskType != "all" { + if taskType != "" { images = filterImagesByTaskType(images, taskType) } images = filterImages(cmd, images) if outputJSON { - return output.WriteSuccessJSON(os.Stdout, output.SuccessEnvelope(map[string]interface{}{"images": images})) - } - printImageTable(images) - return nil -} - -func runAdminImageLs(cmd *cobra.Command, _ []string) error { - client, err := activeAPIClient() - if err != nil { - return err - } - resp, err := client.ListImageRecords(true) - if err != nil { - return cliErrFromAPI(err) - } - images := filterImages(cmd, resp.ImageList) - if outputJSON { - return output.WriteSuccessJSON(os.Stdout, output.SuccessEnvelope(map[string]interface{}{"images": images})) + return output.WriteSuccessJSON(os.Stdout, output.SuccessEnvelope(map[string]interface{}{ + "images": images, + })) } printImageTable(images) return nil @@ -368,784 +118,35 @@ func filterImages(cmd *cobra.Command, images []api.ImageInfo) []api.ImageInfo { return out } -func runImageUpload(cmd *cobra.Command, _ []string) error { - req, err := collectUpload(cmd) - if err != nil { - return err - } - client, err := activeAPIClient() - if err != nil { - return err - } - msg, err := client.UploadImage(req) - return writeImageMessage(msg, err) -} - -func runImageDelete(_ *cobra.Command, args []string) error { - id, err := requiredUintArg(args, "image_label_id", "id") - if err != nil { - return err - } - client, err := activeAPIClient() - if err != nil { - return err - } - msg, err := client.DeleteImage(id) - return writeImageMessage(msg, err) -} - -func runImageDeleteMany(cmd *cobra.Command, _ []string) error { - ids, err := idsFlag(cmd) - if err != nil { - return err - } - client, err := activeAPIClient() - if err != nil { - return err - } - msg, err := client.DeleteImages(ids, false) - return writeImageMessage(msg, err) -} - -func runAdminImageDeleteMany(cmd *cobra.Command, _ []string) error { - ids, err := idsFlag(cmd) - if err != nil { - return err - } - client, err := activeAPIClient() - if err != nil { - return err - } - msg, err := client.DeleteImages(ids, true) - return writeImageMessage(msg, err) -} - -func runImageDescription(cmd *cobra.Command, args []string) error { - id, err := requiredUintArg(args, "image_label_id", "id") - if err != nil { - return err - } - description, _ := cmd.Flags().GetString("description") - if strings.TrimSpace(description) == "" { - return errUsageFromIssues([]usageIssue{missingIssue("description", "image_flag_description")}) - } - client, err := activeAPIClient() - if err != nil { - return err - } - msg, err := client.UpdateImageDescription(api.ImageDescriptionRequest{ID: id, Description: description}, false) - return writeImageMessage(msg, err) -} - -func runAdminImageDescription(cmd *cobra.Command, args []string) error { - id, err := requiredUintArg(args, "image_label_id", "id") - if err != nil { - return err - } - description, _ := cmd.Flags().GetString("description") - if strings.TrimSpace(description) == "" { - return errUsageFromIssues([]usageIssue{missingIssue("description", "image_flag_description")}) - } - client, err := activeAPIClient() - if err != nil { - return err - } - msg, err := client.UpdateImageDescription(api.ImageDescriptionRequest{ID: id, Description: description}, true) - return writeImageMessage(msg, err) -} - -func runImageType(cmd *cobra.Command, args []string) error { - id, err := requiredUintArg(args, "image_label_id", "id") - if err != nil { - return err - } - taskType, _ := cmd.Flags().GetString("type") - if err := validateEnum("type", taskType, imageTaskWriteTypes); err != nil { - return err - } - client, err := activeAPIClient() - if err != nil { - return err - } - msg, err := client.UpdateImageType(api.ImageTypeRequest{ID: id, TaskType: taskType}, false) - return writeImageMessage(msg, err) -} - -func runAdminImageType(cmd *cobra.Command, args []string) error { - id, err := requiredUintArg(args, "image_label_id", "id") - if err != nil { - return err - } - taskType, _ := cmd.Flags().GetString("type") - if err := validateEnum("type", taskType, imageTaskWriteTypes); err != nil { - return err - } - client, err := activeAPIClient() - if err != nil { - return err - } - msg, err := client.UpdateImageType(api.ImageTypeRequest{ID: id, TaskType: taskType}, true) - return writeImageMessage(msg, err) -} - -func runImageTags(cmd *cobra.Command, args []string) error { - id, err := requiredUintArg(args, "image_label_id", "id") - if err != nil { - return err - } - tags := csvFlag(cmd, "tags") - client, err := activeAPIClient() - if err != nil { - return err - } - msg, err := client.UpdateImageTags(api.ImageTagsRequest{ID: id, Tags: tags}, false) - return writeImageMessage(msg, err) -} - -func runAdminImageTags(cmd *cobra.Command, args []string) error { - id, err := requiredUintArg(args, "image_label_id", "id") - if err != nil { - return err - } - tags := csvFlag(cmd, "tags") - client, err := activeAPIClient() - if err != nil { - return err - } - msg, err := client.UpdateImageTags(api.ImageTagsRequest{ID: id, Tags: tags}, true) - return writeImageMessage(msg, err) -} - -func runImageArch(cmd *cobra.Command, args []string) error { - id, err := requiredUintArg(args, "image_label_id", "id") - if err != nil { - return err - } - archs, err := requiredArchitectures(cmd) - if err != nil { - return err - } - client, err := activeAPIClient() - if err != nil { - return err - } - msg, err := client.UpdateImageArch(api.ImageArchRequest{ID: id, Archs: archs}, false) - return writeImageMessage(msg, err) -} - -func runAdminImageArch(cmd *cobra.Command, args []string) error { - id, err := requiredUintArg(args, "image_label_id", "id") - if err != nil { - return err - } - archs, err := requiredArchitectures(cmd) - if err != nil { - return err - } - client, err := activeAPIClient() - if err != nil { - return err +func printImageTable(images []api.ImageInfo) { + fmt.Printf("%s %s %s %s %s %s\n", + i18n.PadRight(i18n.T("table_id"), 8), + i18n.PadRight(i18n.T("table_image"), 48), + i18n.PadRight(i18n.T("table_type"), 12), + i18n.PadRight(i18n.T("table_visibility"), 14), + i18n.PadRight(i18n.T("table_arch"), 24), + i18n.PadRight(i18n.T("table_owner"), 18)) + for _, image := range images { + fmt.Printf("%s %s %s %s %s %s\n", + i18n.PadRight(fmt.Sprintf("%d", image.ID), 8), + i18n.PadRight(image.ImageLink, 48), + i18n.PadRight(image.TaskType, 12), + i18n.PadRight(image.ImageShareStatus, 14), + i18n.PadRight(strings.Join(image.Archs, ","), 24), + i18n.PadRight(image.UserInfo.Nickname, 18)) } - msg, err := client.UpdateImageArch(api.ImageArchRequest{ID: id, Archs: archs}, true) - return writeImageMessage(msg, err) -} - -func runImageShareAdd(cmd *cobra.Command, args []string) error { - imageID, err := requiredUintArg(args, "image_label_id", "image-id") - if err != nil { - return err - } - targets, err := idsFlag(cmd) - if err != nil { - return err - } - shareType, _ := cmd.Flags().GetString("share-type") - if err := validateEnum("share-type", shareType, imageShareTypes); err != nil { - return err - } - client, err := activeAPIClient() - if err != nil { - return err - } - msg, err := client.ShareImage(api.ImageShareRequest{ImageID: imageID, IDList: targets, Type: shareType}) - return writeImageMessage(msg, err) -} - -func runImageShareRemove(cmd *cobra.Command, args []string) error { - imageID, err := requiredUintArg(args, "image_label_id", "image-id") - if err != nil { - return err - } - targetID, _ := cmd.Flags().GetUint("target-id") - if targetID == 0 { - return errUsageFromIssues([]usageIssue{missingIssue("target-id", "image_flag_target-id")}) - } - shareType, _ := cmd.Flags().GetString("share-type") - if err := validateEnum("share-type", shareType, imageShareTypes); err != nil { - return err - } - client, err := activeAPIClient() - if err != nil { - return err - } - msg, err := client.CancelShareImage(api.ImageCancelShareRequest{ImageID: imageID, ID: targetID, Type: shareType}) - return writeImageMessage(msg, err) -} - -func runImageShareLs(_ *cobra.Command, args []string) error { - imageID, err := requiredUintArg(args, "image_label_id", "image-id") - if err != nil { - return err - } - client, err := activeAPIClient() - if err != nil { - return err - } - resp, err := client.GetImageGrants(imageID) - if err != nil { - return cliErrFromAPI(err) - } - if outputJSON { - return output.WriteSuccessJSON(os.Stdout, output.SuccessEnvelope(map[string]interface{}{"grants": resp})) - } - printGrantTable(resp) - return nil -} - -func runImageShareUsers(cmd *cobra.Command, args []string) error { - imageID, err := requiredUintArg(args, "image_label_id", "image-id") - if err != nil { - return err - } - name, _ := cmd.Flags().GetString("name") - client, err := activeAPIClient() - if err != nil { - return err - } - resp, err := client.ListUngrantedUsers(imageID, name) - if err != nil { - return cliErrFromAPI(err) - } - if outputJSON { - return output.WriteSuccessJSON(os.Stdout, output.SuccessEnvelope(map[string]interface{}{"users": resp.UserList})) - } - printUserGrantTable(resp.UserList) - return nil -} - -func runImageShareAccounts(_ *cobra.Command, args []string) error { - imageID, err := requiredUintArg(args, "image_label_id", "image-id") - if err != nil { - return err - } - client, err := activeAPIClient() - if err != nil { - return err - } - resp, err := client.ListUngrantedAccounts(imageID) - if err != nil { - return cliErrFromAPI(err) - } - if outputJSON { - return output.WriteSuccessJSON(os.Stdout, output.SuccessEnvelope(map[string]interface{}{"accounts": resp.AccountList})) - } - printAccountGrantTable(resp.AccountList) - return nil -} - -func runImageValid(cmd *cobra.Command, _ []string) error { - pairs, err := linkPairsFlag(cmd) - if err != nil { - return err - } - client, err := activeAPIClient() - if err != nil { - return err - } - resp, err := client.CheckImageLinks(pairs) - if err != nil { - return cliErrFromAPI(err) - } - if outputJSON { - return output.WriteSuccessJSON(os.Stdout, output.SuccessEnvelope(map[string]interface{}{"invalid_pairs": resp.InvalidPairs})) - } - printLinkPairTable(resp.InvalidPairs) - return nil -} - -func runImageCudaLs(_ *cobra.Command, _ []string) error { - client, err := activeAPIClient() - if err != nil { - return err - } - resp, err := client.ListCudaBaseImages() - if err != nil { - return cliErrFromAPI(err) - } - if outputJSON { - return output.WriteSuccessJSON(os.Stdout, output.SuccessEnvelope(map[string]interface{}{"cuda_base_images": resp.CudaBaseImages})) - } - printCudaTable(resp.CudaBaseImages) - return nil -} - -func runAdminImageCudaAdd(cmd *cobra.Command, _ []string) error { - imageLabel, _ := cmd.Flags().GetString("image-label") - label, _ := cmd.Flags().GetString("label") - value, _ := cmd.Flags().GetString("value") - var issues []usageIssue - if strings.TrimSpace(imageLabel) == "" { - issues = append(issues, missingIssue("image-label", "image_flag_image-label")) - } - if strings.TrimSpace(label) == "" { - issues = append(issues, missingIssue("label", "image_flag_label")) - } - if strings.TrimSpace(value) == "" { - issues = append(issues, missingIssue("value", "image_flag_value")) - } - if len(issues) > 0 { - return errUsageFromIssues(issues) - } - client, err := activeAPIClient() - if err != nil { - return err - } - msg, err := client.AdminAddCudaBaseImage(api.CudaBaseImageRequest{ImageLabel: imageLabel, Label: label, Value: value}) - return writeImageMessage(msg, err) -} - -func runAdminImageCudaDelete(_ *cobra.Command, args []string) error { - id, err := requiredUintArg(args, "image_label_id", "id") - if err != nil { - return err - } - client, err := activeAPIClient() - if err != nil { - return err - } - msg, err := client.AdminDeleteCudaBaseImage(id) - return writeImageMessage(msg, err) -} - -func runImageHarborInfo(_ *cobra.Command, _ []string) error { - client, err := activeAPIClient() - if err != nil { - return err - } - resp, err := client.GetHarbor() - if err != nil { - return cliErrFromAPI(err) - } - if outputJSON { - return output.WriteSuccessJSON(os.Stdout, output.SuccessEnvelope(map[string]interface{}{"harbor": resp})) - } - fmt.Printf("%s: %s\n", i18n.T("image_output_harbor"), resp.IP) - return nil -} - -func runImageHarborCredential(cmd *cobra.Command, _ []string) error { - if err := requireYes(cmd); err != nil { - return err - } - client, err := activeAPIClient() - if err != nil { - return err - } - resp, err := client.GetCredential() - if err != nil { - return cliErrFromAPI(err) - } - if outputJSON { - return output.WriteSuccessJSON(os.Stdout, output.SuccessEnvelope(map[string]interface{}{"credential": resp})) - } - fmt.Printf("%s: %s\n%s: %s\n", i18n.T("image_output_username"), deref(resp.Name), i18n.T("image_output_password"), deref(resp.Password)) - return nil -} - -func runImageQuotaGet(_ *cobra.Command, _ []string) error { - client, err := activeAPIClient() - if err != nil { - return err - } - resp, err := client.GetQuota() - if err != nil { - return cliErrFromAPI(err) - } - if outputJSON { - return output.WriteSuccessJSON(os.Stdout, output.SuccessEnvelope(map[string]interface{}{"quota": resp})) - } - fmt.Printf("%s: %s\n%s: %.2f GB\n%s: %.2f GB\n", i18n.T("image_output_project"), resp.Project, i18n.T("image_output_used"), resp.Used, i18n.T("image_output_quota"), resp.Quota) - return nil -} - -func runImageQuotaSet(cmd *cobra.Command, _ []string) error { - size, _ := cmd.Flags().GetInt64("size") - if size <= 0 { - return errUsageFromIssues([]usageIssue{missingIssue("size", "image_flag_size")}) - } - client, err := activeAPIClient() - if err != nil { - return err - } - msg, err := client.UpdateQuota(size) - return writeImageMessage(msg, err) -} - -func runAdminImagePublic(_ *cobra.Command, args []string) error { - id, err := requiredUintArg(args, "image_label_id", "id") - if err != nil { - return err - } - client, err := activeAPIClient() - if err != nil { - return err - } - msg, err := client.TogglePublic(id) - return writeImageMessage(msg, err) -} - -func collectPipAptBuild(cmd *cobra.Command) (api.PipAptBuildRequest, error) { - name, tag, description, tags, archs, template, err := commonBuildFlags(cmd) - if err != nil { - return api.PipAptBuildRequest{}, err - } - image, _ := cmd.Flags().GetString("image") - if strings.TrimSpace(image) == "" { - return api.PipAptBuildRequest{}, errUsageFromIssues([]usageIssue{missingIssue("image", "image_flag_image")}) - } - requirements, _ := cmd.Flags().GetString("requirements") - packages, _ := cmd.Flags().GetString("packages") - return api.PipAptBuildRequest{Image: image, Requirements: requirements, Packages: packages, Description: description, Name: name, Tag: tag, Tags: tags, Template: template, Archs: archs}, nil -} - -func collectDockerfileBuild(cmd *cobra.Command) (api.DockerfileBuildRequest, error) { - name, tag, description, tags, archs, template, err := commonBuildFlags(cmd) - if err != nil { - return api.DockerfileBuildRequest{}, err - } - dockerfile, err := contentFromFlags(cmd, "dockerfile", "file") - if err != nil { - return api.DockerfileBuildRequest{}, err - } - return api.DockerfileBuildRequest{Dockerfile: dockerfile, Description: description, Name: name, Tag: tag, Tags: tags, Template: template, Archs: archs}, nil -} - -func collectEnvdBuild(cmd *cobra.Command) (api.EnvdBuildRequest, error) { - name, tag, description, tags, archs, template, err := commonBuildFlags(cmd) - if err != nil { - return api.EnvdBuildRequest{}, err - } - envd, err := contentFromFlags(cmd, "envd", "file") - if err != nil { - return api.EnvdBuildRequest{}, err - } - python, _ := cmd.Flags().GetString("python") - base, _ := cmd.Flags().GetString("base") - buildSource, _ := cmd.Flags().GetString("build-source") - if buildSource == "" { - buildSource = "EnvdAdvanced" - } - if err := validateEnum("build-source", buildSource, imageBuildSources); err != nil { - return api.EnvdBuildRequest{}, err - } - return api.EnvdBuildRequest{Envd: envd, Description: description, Name: name, Tag: tag, Python: python, Base: base, Tags: tags, Template: template, BuildSource: buildSource, Archs: archs}, nil -} - -func commonBuildFlags(cmd *cobra.Command) (string, string, string, []string, []string, string, error) { - name, _ := cmd.Flags().GetString("name") - tag, _ := cmd.Flags().GetString("tag") - description, _ := cmd.Flags().GetString("description") - template, _ := cmd.Flags().GetString("template") - tags := csvFlag(cmd, "tags") - archs := csvFlag(cmd, "archs") - if len(archs) == 0 { - archs = []string{"linux/amd64"} - } - var issues []usageIssue - if strings.TrimSpace(name) == "" { - issues = append(issues, missingIssue("name", "image_flag_name")) - } - if strings.TrimSpace(tag) == "" { - issues = append(issues, missingIssue("tag", "image_flag_tag")) - } - if len(issues) > 0 { - return "", "", "", nil, nil, "", errUsageFromIssues(issues) - } - return name, tag, description, tags, archs, template, nil -} - -func collectUpload(cmd *cobra.Command) (api.ImageUploadRequest, error) { - image, _ := cmd.Flags().GetString("image") - taskType, _ := cmd.Flags().GetString("type") - description, _ := cmd.Flags().GetString("description") - if taskType == "" { - taskType = "custom" - } - var issues []usageIssue - if strings.TrimSpace(image) == "" { - issues = append(issues, missingIssue("image", "image_flag_image")) - } - if taskType != "" && !slices.Contains(imageTaskWriteTypes, taskType) { - issues = append(issues, invalidIssue("type", i18n.T("err_invalid_image_type", taskType))) - } - if len(issues) > 0 { - return api.ImageUploadRequest{}, errUsageFromIssues(issues) - } - return api.ImageUploadRequest{ImageLink: image, TaskType: taskType, Description: description, Tags: csvFlag(cmd, "tags"), Archs: csvFlag(cmd, "archs")}, nil -} - -func contentFromFlags(cmd *cobra.Command, directFlag, fileFlag string) (string, error) { - direct, _ := cmd.Flags().GetString(directFlag) - file, _ := cmd.Flags().GetString(fileFlag) - if cmd.Flags().Changed(directFlag) && cmd.Flags().Changed(fileFlag) { - return "", errUsageFromIssues([]usageIssue{invalidIssue(directFlag, i18n.T("err_mutually_exclusive_flags", directFlag, fileFlag))}) - } - if file != "" { - b, err := os.ReadFile(file) - if err != nil { - return "", err - } - direct = string(b) - } - if strings.TrimSpace(direct) == "" { - return "", errUsageFromIssues([]usageIssue{missingIssue(directFlag, "image_flag_"+directFlag)}) - } - return direct, nil -} - -func requiredArchitectures(cmd *cobra.Command) ([]string, error) { - archs := csvFlag(cmd, "archs") - if len(archs) == 0 { - return nil, errUsageFromIssues([]usageIssue{missingIssue("archs", "image_flag_archs")}) - } - for _, arch := range archs { - if !slices.Contains(imageArchitectures, arch) { - return nil, errUsageFromIssues([]usageIssue{invalidIssue("archs", i18n.T("err_invalid_image_arch", arch))}) - } - } - return archs, nil -} - -func idsFlag(cmd *cobra.Command) ([]uint, error) { - raw, _ := cmd.Flags().GetString("ids") - ids, err := parseIDs(raw) - if err != nil { - return nil, err - } - if len(ids) == 0 { - return nil, errUsageFromIssues([]usageIssue{missingIssue("ids", "image_flag_ids")}) - } - return ids, nil -} - -func parseIDs(raw string) ([]uint, error) { - parts := splitCSV(raw) - out := make([]uint, 0, len(parts)) - for _, part := range parts { - v, err := strconv.ParseUint(part, 10, 0) - if err != nil { - return nil, errUsageFromIssues([]usageIssue{{Code: errorcodes.ErrInvalidFlagValue, Message: i18n.T("err_invalid_ids", raw), Field: "ids"}}) - } - out = append(out, uint(v)) - } - return out, nil -} - -func csvFlag(cmd *cobra.Command, name string) []string { - raw, _ := cmd.Flags().GetString(name) - values := splitCSV(raw) - if len(values) == 0 { - return nil - } - return values -} - -func validateEnum(field, value string, allowed []string) error { - if strings.TrimSpace(value) == "" { - return errUsageFromIssues([]usageIssue{missingIssue(field, "image_flag_"+field)}) - } - if !slices.Contains(allowed, value) { - return errUsageFromIssues([]usageIssue{invalidIssue(field, i18n.T("err_invalid_image_value", field, value))}) - } - return nil -} - -func linkPairsFlag(cmd *cobra.Command) ([]api.ImageInfoLinkPair, error) { - raw, _ := cmd.Flags().GetString("links") - parts := splitCSV(raw) - if len(parts) == 0 { - return nil, errUsageFromIssues([]usageIssue{missingIssue("links", "image_flag_links")}) - } - pairs := make([]api.ImageInfoLinkPair, 0, len(parts)) - for _, part := range parts { - pairs = append(pairs, api.ImageInfoLinkPair{ImageLink: part}) - } - return pairs, nil -} - -func requireYes(cmd *cobra.Command) error { - yes, _ := cmd.Flags().GetBool("yes") - if !yes { - return errUsageFromIssues([]usageIssue{missingIssue("yes", "flag_yes")}) - } - return nil -} - -func writeImageMessage(msg string, err error) error { - if err != nil { - return cliErrFromAPI(err) - } - if outputJSON { - return output.WriteSuccessJSON(os.Stdout, output.SuccessEnvelope(map[string]interface{}{"message": msg})) - } - fmt.Println(i18n.T("image_success", emptyDash(msg))) - return nil -} - -func printKanikoTable(items []api.KanikoInfo) { - fmt.Printf("%s %s %s %s %s %s\n", i18n.PadRight(i18n.T("table_id"), 8), i18n.PadRight(i18n.T("table_name"), 28), i18n.PadRight(i18n.T("table_status"), 12), i18n.PadRight(i18n.T("table_image"), 42), i18n.PadRight(i18n.T("table_owner"), 18), i18n.PadRight(i18n.T("table_created"), 22)) - for _, item := range items { - fmt.Printf("%s %s %s %s %s %s\n", i18n.PadRight(strconv.FormatUint(uint64(item.ID), 10), 8), i18n.PadRight(item.ImagePackName, 28), i18n.PadRight(item.Status, 12), i18n.PadRight(item.ImageLink, 42), i18n.PadRight(item.UserInfo.Nickname, 18), i18n.PadRight(item.CreatedAt.Format("2006-01-02 15:04:05"), 22)) - } -} - -func printImageTable(items []api.ImageInfo) { - fmt.Printf("%s %s %s %s %s %s\n", i18n.PadRight(i18n.T("table_id"), 8), i18n.PadRight(i18n.T("table_image"), 48), i18n.PadRight(i18n.T("table_type"), 12), i18n.PadRight(i18n.T("table_visibility"), 14), i18n.PadRight(i18n.T("table_arch"), 24), i18n.PadRight(i18n.T("table_owner"), 18)) - for _, item := range items { - fmt.Printf("%s %s %s %s %s %s\n", i18n.PadRight(strconv.FormatUint(uint64(item.ID), 10), 8), i18n.PadRight(item.ImageLink, 48), i18n.PadRight(item.TaskType, 12), i18n.PadRight(item.ImageShareStatus, 14), i18n.PadRight(strings.Join(item.Archs, ","), 24), i18n.PadRight(item.UserInfo.Nickname, 18)) - } -} - -func printCudaTable(items []api.CudaBaseImage) { - fmt.Printf("%s %s %s %s\n", i18n.PadRight(i18n.T("table_id"), 8), i18n.PadRight(i18n.T("table_label"), 24), i18n.PadRight(i18n.T("table_image_label"), 24), i18n.PadRight(i18n.T("table_value"), 48)) - for _, item := range items { - fmt.Printf("%s %s %s %s\n", i18n.PadRight(strconv.FormatUint(uint64(item.ID), 10), 8), i18n.PadRight(item.Label, 24), i18n.PadRight(item.ImageLabel, 24), i18n.PadRight(item.Value, 48)) - } -} - -func printGrantTable(resp *api.ImageGrantResponse) { - if resp == nil { - return - } - printUserGrantTable(resp.UserList) - printAccountGrantTable(resp.AccountList) -} - -func printUserGrantTable(items []api.ImageGrantedUser) { - fmt.Printf("%s %s %s\n", i18n.PadRight(i18n.T("table_id"), 8), i18n.PadRight(i18n.T("table_username"), 24), i18n.PadRight(i18n.T("table_nickname"), 24)) - for _, item := range items { - fmt.Printf("%s %s %s\n", i18n.PadRight(strconv.FormatUint(uint64(item.ID), 10), 8), i18n.PadRight(item.Name, 24), i18n.PadRight(item.Nickname, 24)) - } -} - -func printAccountGrantTable(items []api.ImageGrantedAccount) { - fmt.Printf("%s %s\n", i18n.PadRight(i18n.T("table_id"), 8), i18n.PadRight(i18n.T("table_name"), 24)) - for _, item := range items { - fmt.Printf("%s %s\n", i18n.PadRight(strconv.FormatUint(uint64(item.ID), 10), 8), i18n.PadRight(item.Name, 24)) - } -} - -func printLinkPairTable(items []api.ImageInfoLinkPair) { - fmt.Printf("%s\n", i18n.PadRight(i18n.T("table_image"), 64)) - for _, item := range items { - fmt.Printf("%s\n", i18n.PadRight(item.ImageLink, 64)) - } -} - -func deref(v *string) string { - if v == nil { - return "" - } - return *v } func init() { - addCommonBuildFlags(imageBuildPipAptCmd) - imageBuildPipAptCmd.Flags().String("image", "", "Base image") - imageBuildPipAptCmd.Flags().String("packages", "", "APT packages") - imageBuildPipAptCmd.Flags().String("requirements", "", "Python requirements") - addCommonBuildFlags(imageBuildDockerfileCmd) - imageBuildDockerfileCmd.Flags().String("dockerfile", "", "Dockerfile content") - imageBuildDockerfileCmd.Flags().String("file", "", "Read Dockerfile from file") - addCommonBuildFlags(imageBuildEnvdCmd) - imageBuildEnvdCmd.Flags().String("envd", "", "envd content") - imageBuildEnvdCmd.Flags().String("file", "", "Read envd from file") - imageBuildEnvdCmd.Flags().String("python", "", "Python version") - imageBuildEnvdCmd.Flags().String("base", "", "envd base image") - imageBuildEnvdCmd.Flags().String("build-source", "EnvdAdvanced", "envd build source") - imageBuildRemoveCmd.Flags().String("ids", "", "Comma-separated IDs") - imageBuildCmd.AddCommand(imageBuildLsCmd, imageBuildPipAptCmd, imageBuildDockerfileCmd, imageBuildEnvdCmd, imageBuildRemoveCmd, imageBuildGetCmd, imageBuildTemplateCmd, imageBuildPodCmd) - imageLsCmd.Flags().Bool("available", false, "List images available for creating jobs") imageLsCmd.Flags().String("type", "", "Filter by job type") imageLsCmd.Flags().String("arch", "", "Filter by image architecture") imageLsCmd.Flags().String("visibility", "", "Filter by image visibility") imageLsCmd.Flags().String("owner", "", "Filter by owner username or nickname") imageLsCmd.Flags().String("search", "", "Filter by image link or description substring") - adminImageLsCmd.Flags().String("arch", "", "Filter by image architecture") - adminImageLsCmd.Flags().String("visibility", "", "Filter by image visibility") - adminImageLsCmd.Flags().String("owner", "", "Filter by owner username or nickname") - adminImageLsCmd.Flags().String("search", "", "Filter by image link or description substring") - imageUploadCmd.Flags().String("image", "", "Image link") - imageUploadCmd.Flags().String("type", "custom", "Image task type") - imageUploadCmd.Flags().String("description", "", "Description") - imageUploadCmd.Flags().String("tags", "", "Comma-separated tags") - imageUploadCmd.Flags().String("archs", "", "Comma-separated architectures") - imageDeleteManyCmd.Flags().String("ids", "", "Comma-separated IDs") - imageDescriptionCmd.Flags().String("description", "", "Description") - imageTypeCmd.Flags().String("type", "", "Image task type") - imageTagsCmd.Flags().String("tags", "", "Comma-separated tags") - imageArchCmd.Flags().String("archs", "", "Comma-separated architectures") - imageValidCmd.Flags().String("links", "", "Comma-separated image links") - - imageShareAddCmd.Flags().String("ids", "", "Comma-separated target IDs") - imageShareAddCmd.Flags().String("share-type", "user", "Share type") - imageShareRemoveCmd.Flags().Uint("target-id", 0, "Share target ID") - imageShareRemoveCmd.Flags().String("share-type", "user", "Share type") - imageShareUsersCmd.Flags().String("name", "", "Filter by username or nickname") - imageShareCmd.AddCommand(imageShareLsCmd, imageShareUsersCmd, imageShareAccountsCmd, imageShareAddCmd, imageShareRemoveCmd) - - imageCudaCmd.AddCommand(imageCudaLsCmd) - - imageHarborCredentialCmd.Flags().BoolP("yes", "y", false, "Confirm sensitive operation") - imageHarborCmd.AddCommand(imageHarborInfoCmd, imageHarborCredentialCmd) - imageQuotaSetCmd.Flags().Int64("size", 0, "Quota size") - imageQuotaCmd.AddCommand(imageQuotaGetCmd, imageQuotaSetCmd) - - adminImageBuildRemoveCmd.Flags().String("ids", "", "Comma-separated IDs") - adminImageDeleteManyCmd.Flags().String("ids", "", "Comma-separated IDs") - adminImageDescriptionCmd.Flags().String("description", "", "Description") - adminImageTypeCmd.Flags().String("type", "", "Image task type") - adminImageTagsCmd.Flags().String("tags", "", "Comma-separated tags") - adminImageArchCmd.Flags().String("archs", "", "Comma-separated architectures") - adminImageCudaAddCmd.Flags().String("image-label", "", "Image label") - adminImageCudaAddCmd.Flags().String("label", "", "Display label") - adminImageCudaAddCmd.Flags().String("value", "", "Image value") - - completion.RegisterFlagValue([]string{"image", "ls"}, "type", staticValueCompleter(imageTaskFilterTypes, nil)) - completion.RegisterFlagValue([]string{"image", "ls"}, "arch", staticValueCompleter(imageArchitectures, nil)) - completion.RegisterFlagValue([]string{"image", "arch"}, "archs", staticValueCompleter(imageArchitectures, nil)) - completion.RegisterFlagValue([]string{"admin", "image", "arch"}, "archs", staticValueCompleter(imageArchitectures, nil)) + completion.RegisterFlagValue([]string{"image", "ls"}, "type", staticValueCompleter(imageTaskTypes, nil)) + completion.RegisterFlagValue([]string{"image", "ls"}, "arch", staticValueCompleter([]string{"linux/amd64", "linux/arm64"}, nil)) completion.RegisterFlagValue([]string{"image", "ls"}, "visibility", staticValueCompleter(imageVisibilityTypes, nil)) - completion.RegisterFlagValue([]string{"image", "upload"}, "type", staticValueCompleter(imageTaskWriteTypes, nil)) - completion.RegisterFlagValue([]string{"image", "type"}, "type", staticValueCompleter(imageTaskWriteTypes, nil)) - completion.RegisterFlagValue([]string{"admin", "image", "type"}, "type", staticValueCompleter(imageTaskWriteTypes, nil)) - completion.RegisterFlagValue([]string{"image", "share", "add"}, "share-type", staticValueCompleter(imageShareTypes, nil)) - completion.RegisterFlagValue([]string{"image", "share", "remove"}, "share-type", staticValueCompleter(imageShareTypes, nil)) - completion.RegisterFlagValue([]string{"image", "build", "envd"}, "build-source", staticValueCompleter(imageBuildSources, nil)) - - imageCmd.AddCommand(imageBuildCmd, imageLsCmd, imageUploadCmd, imageDeleteCmd, imageDeleteManyCmd, imageDescriptionCmd, imageTypeCmd, imageTagsCmd, imageArchCmd, imageValidCmd, imageShareCmd, imageCudaCmd, imageHarborCmd, imageQuotaCmd) + imageCmd.AddCommand(imageLsCmd) rootCmd.AddCommand(imageCmd) - adminImageCudaCmd.AddCommand(adminImageCudaAddCmd, adminImageCudaDeleteCmd) - adminImageCmd.AddCommand(adminImageBuildLsCmd, adminImageBuildRemoveCmd, adminImageLsCmd, adminImageDeleteManyCmd, adminImageDescriptionCmd, adminImageTypeCmd, adminImageTagsCmd, adminImageArchCmd, adminImagePublicCmd, adminImageCudaCmd) - adminCmd.AddCommand(adminImageCmd) -} - -func addCommonBuildFlags(cmd *cobra.Command) { - cmd.Flags().String("name", "", "Image name") - cmd.Flags().String("tag", "", "Image tag") - cmd.Flags().String("description", "", "Description") - cmd.Flags().String("tags", "", "Comma-separated tags") - cmd.Flags().String("archs", "", "Comma-separated architectures") - cmd.Flags().String("template", "", "Template text") } diff --git a/cli/cmd/image_test.go b/cli/cmd/image_test.go deleted file mode 100644 index 49df8397f..000000000 --- a/cli/cmd/image_test.go +++ /dev/null @@ -1,83 +0,0 @@ -package cmd - -import ( - "io" - "os" - "strings" - "testing" - - "github.com/raids-lab/crater/cli/internal/api" - "github.com/raids-lab/crater/cli/internal/i18n" - "github.com/spf13/cobra" -) - -func TestPrintImageTableIncludesVisibility(t *testing.T) { - previousLanguage := i18n.GetCurrentLanguage() - i18n.SetLanguage("en") - t.Cleanup(func() { i18n.SetLanguage(previousLanguage) }) - - got := captureImageTestStdout(t, func() { - printImageTable([]api.ImageInfo{{ - ID: 1, - ImageLink: "registry.example/demo:v1", - TaskType: "custom", - ImageShareStatus: "Private", - Archs: []string{"linux/amd64"}, - UserInfo: api.UserInfo{Nickname: "alice"}, - }}) - }) - if !strings.Contains(got, "VISIBILITY") || !strings.Contains(got, "Private") { - t.Fatalf("table must display image visibility, got %q", got) - } - if strings.Contains(got, "CREATED") { - t.Fatalf("table must not display the removed CREATED column, got %q", got) - } -} - -func TestImageCommandGroupsRejectUnknownSubcommands(t *testing.T) { - tests := []struct { - name string - cmd *cobra.Command - }{ - {name: "build", cmd: imageBuildCmd}, - {name: "share", cmd: imageShareCmd}, - {name: "user cuda", cmd: imageCudaCmd}, - {name: "harbor", cmd: imageHarborCmd}, - {name: "quota", cmd: imageQuotaCmd}, - {name: "admin image", cmd: adminImageCmd}, - {name: "admin cuda", cmd: adminImageCudaCmd}, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - // Cobra treats arguments on a command group as positional input when - // no matching child exists. Exercise the group's RunE directly so the - // contract remains independent of the process-wide root command state. - err := tt.cmd.RunE(tt.cmd, []string{"unknown"}) - if err == nil || !strings.Contains(err.Error(), `unknown command "unknown"`) { - t.Fatalf("unknown subcommand error = %v", err) - } - }) - } -} - -func captureImageTestStdout(t *testing.T, fn func()) string { - t.Helper() - old := os.Stdout - r, w, err := os.Pipe() - if err != nil { - t.Fatalf("create stdout pipe: %v", err) - } - os.Stdout = w - t.Cleanup(func() { os.Stdout = old }) - - fn() - if err := w.Close(); err != nil { - t.Fatalf("close stdout pipe: %v", err) - } - output, err := io.ReadAll(r) - if err != nil { - t.Fatalf("read stdout pipe: %v", err) - } - return string(output) -} diff --git a/cli/cmd/root.go b/cli/cmd/root.go index d4c5b76b0..9e05a7085 100644 --- a/cli/cmd/root.go +++ b/cli/cmd/root.go @@ -210,7 +210,14 @@ func updateAllCommands(cmd *cobra.Command) { // 3. Update Flags for this command cmd.Flags().VisitAll(func(f *pflag.Flag) { - if usage, ok := translatedFlagUsage(keyPath, f.Name); ok { + commandFlagKey := keyPath + "_flag_" + f.Name + if usage := i18n.T(commandFlagKey); usage != commandFlagKey { + f.Usage = usage + return + } + flagKey := "flag_" + f.Name + usage := i18n.T(flagKey) + if usage != flagKey { f.Usage = usage } }) @@ -220,22 +227,3 @@ func updateAllCommands(cmd *cobra.Command) { updateAllCommands(sub) } } - -func translatedFlagUsage(keyPath, flagName string) (string, bool) { - keys := []string{keyPath + "_flag_" + flagName} - if isImageCommandPath(keyPath) { - keys = append(keys, "image_flag_"+flagName) - } - keys = append(keys, "flag_"+flagName) - for _, key := range keys { - if usage := i18n.T(key); usage != key { - return usage, true - } - } - return "", false -} - -func isImageCommandPath(keyPath string) bool { - return keyPath == "image" || strings.HasPrefix(keyPath, "image_") || - keyPath == "admin_image" || strings.HasPrefix(keyPath, "admin_image_") -} diff --git a/cli/cmd/root_test.go b/cli/cmd/root_test.go index 2e17f1f61..46b65a4b8 100644 --- a/cli/cmd/root_test.go +++ b/cli/cmd/root_test.go @@ -3,8 +3,6 @@ package cmd import ( "os" "testing" - - "github.com/raids-lab/crater/cli/internal/i18n" ) func TestBootstrapJSONFlagFromArgsMatchesPflagBoolSemantics(t *testing.T) { @@ -60,53 +58,3 @@ func TestBootstrapJSONFlagFromArgsMatchesPflagBoolSemantics(t *testing.T) { }) } } - -func TestTranslatedFlagUsagePrecedence(t *testing.T) { - previousLanguage := i18n.GetCurrentLanguage() - i18n.SetLanguage("en") - t.Cleanup(func() { i18n.SetLanguage(previousLanguage) }) - - tests := []struct { - name string - keyPath string - flagName string - want string - }{ - { - name: "command-specific text wins", - keyPath: "order_submit", - flagName: "name", - want: "Approval target name", - }, - { - name: "image-domain text wins over global text", - keyPath: "image_upload", - flagName: "type", - want: "Image task type", - }, - { - name: "admin image uses image-domain text", - keyPath: "admin_image_type", - flagName: "type", - want: "Image task type", - }, - { - name: "non-image command falls back to global text", - keyPath: "resource_ls", - flagName: "type", - want: "Filter by type", - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - got, ok := translatedFlagUsage(tt.keyPath, tt.flagName) - if !ok { - t.Fatalf("translatedFlagUsage(%q, %q) did not find a translation", tt.keyPath, tt.flagName) - } - if got != tt.want { - t.Fatalf("translatedFlagUsage(%q, %q) = %q, want %q", tt.keyPath, tt.flagName, got, tt.want) - } - }) - } -} diff --git a/cli/docs/COMMANDS.md b/cli/docs/COMMANDS.md index 8ce95431c..d573de0b9 100644 --- a/cli/docs/COMMANDS.md +++ b/cli/docs/COMMANDS.md @@ -513,7 +513,7 @@ ## 7. 镜像模块 (image) -本模块提供容器镜像、镜像构建、分享、CUDA base image 和 Harbor 项目管理能力。所有命令均要求已有 active credentials。用户可操作资源使用 `crater image ...`；管理员/平台级资源统一使用 `crater admin image ...`，不得使用 `--admin` 切换。 +本模块提供容器镜像信息的只读查询能力。所有命令均要求已有 active credentials。 ### `crater image ls` - **描述**: 列出当前账号可见的镜像。 @@ -532,55 +532,6 @@ - **`--json` 的 `data`**：`images`（数组，元素与平台镜像响应一致，过滤后返回）。 - **状态**: [x] Completed -### Image Build Commands -- `crater image build ls`: `/api/v1/images/kaniko` -- `crater image build get `: `/api/v1/images/getbyname?name=...` -- `crater image build template `: `/api/v1/images/template?name=...` -- `crater image build pod `: `/api/v1/images/podname?id=...` -- `crater image build pip-apt --name NAME --tag TAG --image BASE [--packages TEXT] [--requirements TEXT]` -- `crater image build dockerfile --name NAME --tag TAG (--dockerfile TEXT | --file PATH)` -- `crater image build envd --name NAME --tag TAG (--envd TEXT | --file PATH) [--build-source EnvdAdvanced|EnvdRaw]` -- `crater image build remove --ids 1,2` -- Admin variants: - - `crater admin image build-ls` - - `crater admin image build-remove --ids 1,2` -- JSON payload keys: `builds`, `build`, `template`, `pod`, `message`. - -### Image Record Commands -- `crater image upload --image IMAGE [--type jupyter|webide|custom|pytorch|tensorflow]` -- `crater image delete ` -- `crater image delete-many --ids 1,2` -- `crater image description --description TEXT` -- `crater image type --type jupyter|webide|custom|pytorch|tensorflow` -- `crater image tags --tags a,b` -- `crater image arch --archs linux/amd64,linux/arm64` -- `crater image valid --links image-a,image-b` -- Admin variants: - - `crater admin image ls` - - `crater admin image delete-many --ids 1,2` - - `crater admin image description --description TEXT` - - `crater admin image type --type jupyter|webide|custom|pytorch|tensorflow` - - `crater admin image tags --tags a,b` - - `crater admin image arch --archs linux/amd64` - - `crater admin image public ` -- `type=all` is accepted only as a local list filter, not as a writable image task type. -- JSON payload keys: `images`, `message`, `invalid_pairs`. - -### Image Share, CUDA, Harbor, And Quota Commands -- `crater image share ls `: `/api/v1/images/share?imageID=...` -- `crater image share users [--name NAME]`: `/api/v1/images/user` -- `crater image share accounts `: `/api/v1/images/account` -- `crater image share add --share-type user|account --ids 1,2` -- `crater image share remove --share-type user|account --target-id ID` -- `crater image cuda ls` -- `crater admin image cuda add --image-label LABEL --label TEXT --value IMAGE` -- `crater admin image cuda delete ` -- `crater image harbor info` -- `crater image harbor credential --yes` -- `crater image quota get|set --size BYTES` -- Harbor credential output contains sensitive data and requires explicit `--yes` in every mode. -- JSON payload keys: `grants`, `users`, `accounts`, `cuda_base_images`, `harbor`, `credential`, `quota`, `message`. - --- ## 7. Additional Read Modules diff --git a/cli/internal/api/image.go b/cli/internal/api/image.go index 9b8b10aa2..4dd734716 100644 --- a/cli/internal/api/image.go +++ b/cli/internal/api/image.go @@ -1,89 +1,27 @@ package api -import ( - "fmt" - "time" -) +import "time" type ImageClient interface { - ListKaniko(admin bool) (*ListKanikoResponse, error) - GetKanikoByName(name string) (*KanikoDetailResponse, error) - GetKanikoTemplateByName(name string) (string, error) - GetKanikoPod(id uint) (*KanikoPodResponse, error) - CreatePipApt(req PipAptBuildRequest) (string, error) - CreateDockerfile(req DockerfileBuildRequest) (string, error) - CreateEnvd(req EnvdBuildRequest) (string, error) - RemoveKaniko(ids []uint, admin bool) (string, error) - ListImageRecords(admin bool) (*ListImageResponse, error) - ListAvailableImages() ([]ImageInfo, error) - UploadImage(req ImageUploadRequest) (string, error) - DeleteImage(id uint) (string, error) - DeleteImages(ids []uint, admin bool) (string, error) - UpdateImageDescription(req ImageDescriptionRequest, admin bool) (string, error) - UpdateImageType(req ImageTypeRequest, admin bool) (string, error) - UpdateImageTags(req ImageTagsRequest, admin bool) (string, error) - UpdateImageArch(req ImageArchRequest, admin bool) (string, error) - TogglePublic(id uint) (string, error) - ShareImage(req ImageShareRequest) (string, error) - CancelShareImage(req ImageCancelShareRequest) (string, error) - GetImageGrants(imageID uint) (*ImageGrantResponse, error) - ListUngrantedUsers(imageID uint, name string) (*ImageGrantResponse, error) - ListUngrantedAccounts(imageID uint) (*ImageGrantResponse, error) - CheckImageLinks(pairs []ImageInfoLinkPair) (*CheckLinkValidityResponse, error) - GetHarbor() (*HarborResponse, error) - GetCredential() (*ProjectCredentialResponse, error) - GetQuota() (*ProjectDetailResponse, error) - UpdateQuota(size int64) (string, error) - ListCudaBaseImages() (*CudaBaseImagesResponse, error) - AdminAddCudaBaseImage(req CudaBaseImageRequest) (string, error) - AdminDeleteCudaBaseImage(id uint) (string, error) + ListImages(available bool) ([]ImageInfo, error) } -type KanikoInfo struct { - ID uint `json:"ID"` - ImageLink string `json:"imageLink"` - Status string `json:"status"` - BuildSource string `json:"buildSource"` - CreatedAt time.Time `json:"createdAt"` - Size int64 `json:"size"` - Description string `json:"description"` - UserInfo UserInfo `json:"userInfo"` - Tags []string `json:"tags"` - ImagePackName string `json:"imagepackName"` - Archs []string `json:"archs"` -} - -type ListKanikoResponse struct { - KanikoList []KanikoInfo `json:"kanikoList"` -} - -type KanikoDetailResponse struct { - ID uint `json:"ID"` - ImageLink string `json:"imageLink"` - Status string `json:"status"` - BuildSource string `json:"buildSource"` - CreatedAt time.Time `json:"createdAt"` - ImagePackName string `json:"imagepackName"` - Description string `json:"description"` - Dockerfile string `json:"dockerfile"` - PodName string `json:"podName"` - PodNameSpace string `json:"podNameSpace"` - NodeName string `json:"nodeName"` +type ImageListResp struct { + ImageList []ImageInfo `json:"imageList"` } -type KanikoPodResponse struct { - Name string `json:"name"` - Namespace string `json:"namespace"` - NodeName string `json:"nodeName"` +type AvailableImageListResp struct { + Images []ImageInfo `json:"images"` } type ImageInfo struct { ID uint `json:"ID"` ImageLink string `json:"imageLink"` Description *string `json:"description"` + Status string `json:"status"` CreatedAt time.Time `json:"createdAt"` - TaskType string `json:"taskType"` IsPublic bool `json:"isPublic"` + TaskType string `json:"taskType"` UserInfo UserInfo `json:"userInfo"` Tags []string `json:"tags"` ImageBuildSource int `json:"imageBuildSource"` @@ -92,440 +30,32 @@ type ImageInfo struct { Archs []string `json:"archs"` } -type ListImageResponse struct { - ImageList []ImageInfo `json:"imageList"` -} - -type AvailableImageListResp struct { - Images []ImageInfo `json:"images"` -} - -type PipAptBuildRequest struct { - Image string `json:"image"` - Requirements string `json:"requirements"` - Packages string `json:"packages"` - Description string `json:"description"` - Name string `json:"name"` - Tag string `json:"tag"` - Tags []string `json:"tags"` - Template string `json:"template"` - Archs []string `json:"archs"` -} - -type DockerfileBuildRequest struct { - Description string `json:"description"` - Dockerfile string `json:"dockerfile"` - Name string `json:"name"` - Tag string `json:"tag"` - Tags []string `json:"tags"` - Template string `json:"template"` - Archs []string `json:"archs"` -} - -type EnvdBuildRequest struct { - Description string `json:"description"` - Envd string `json:"envd"` - Name string `json:"name"` - Tag string `json:"tag"` - Python string `json:"python"` - Base string `json:"base"` - Tags []string `json:"tags"` - Template string `json:"template"` - BuildSource string `json:"buildSource"` - Archs []string `json:"archs"` -} - -type IDListRequest struct { - IDList []uint `json:"idList"` -} - -type ImageUploadRequest struct { - ImageLink string `json:"imageLink"` - TaskType string `json:"taskType"` - Description string `json:"description"` - Tags []string `json:"tags"` - Archs []string `json:"archs"` -} - -type ImageDescriptionRequest struct { - ID uint `json:"id"` - Description string `json:"description"` -} - -type ImageTypeRequest struct { - ID uint `json:"id"` - TaskType string `json:"taskType"` -} - -type ImageTagsRequest struct { - ID uint `json:"id"` - Tags []string `json:"tags"` -} - -type ImageArchRequest struct { - ID uint `json:"id"` - Archs []string `json:"archs"` -} - -type ImageShareRequest struct { - IDList []uint `json:"idList"` - ImageID uint `json:"imageID"` - Type string `json:"type"` -} - -type ImageCancelShareRequest struct { - ID uint `json:"id"` - ImageID uint `json:"imageID"` - Type string `json:"type"` -} - -type ImageInfoLinkPair struct { - ImageLink string `json:"imageLink"` -} - -type CheckLinkValidityRequest struct { - LinkPairs []ImageInfoLinkPair `json:"linkPairs"` -} - -type CheckLinkValidityResponse struct { - InvalidPairs []ImageInfoLinkPair `json:"linkPairs"` -} - -type ImageGrantedUser struct { - ID uint `json:"id"` - Name string `json:"name"` - Nickname string `json:"nickname"` -} - -type ImageGrantedAccount struct { - ID uint `json:"id"` - Name string `json:"name"` -} - -type ImageGrantResponse struct { - UserList []ImageGrantedUser `json:"userList"` - AccountList []ImageGrantedAccount `json:"accountList"` -} - -type HarborResponse struct { - IP string `json:"ip"` -} - -type ProjectCredentialResponse struct { - Name *string `json:"name"` - Password *string `json:"password"` -} - -type ProjectDetailResponse struct { - Used float64 `json:"used"` - Quota float64 `json:"quota"` - Project string `json:"project"` - Total int64 `json:"total"` -} - -type CudaBaseImage struct { - ID uint `json:"id"` - Label string `json:"label"` - ImageLabel string `json:"imageLabel"` - Value string `json:"value"` -} - -type CudaBaseImagesResponse struct { - CudaBaseImages []CudaBaseImage `json:"cudaBaseImages"` -} - -type CudaBaseImageRequest struct { - ImageLabel string `json:"imageLabel"` - Label string `json:"label"` - Value string `json:"value"` -} - -func (c *Client) ListKaniko(admin bool) (*ListKanikoResponse, error) { - path := ImagesPrefix + "/kaniko" - if admin { - path = AdminImagesPrefix + "/kaniko" - } - var result Response[ListKanikoResponse] - if err := c.get(path, nil, &result); err != nil { - return nil, err - } - return &result.Data, nil -} - -func (c *Client) GetKanikoByName(name string) (*KanikoDetailResponse, error) { - var result Response[KanikoDetailResponse] - if err := c.get(ImagesPrefix+"/getbyname", map[string]string{"name": name}, &result); err != nil { - return nil, err - } - return &result.Data, nil -} - -func (c *Client) GetKanikoTemplateByName(name string) (string, error) { - var result Response[string] - if err := c.get(ImagesPrefix+"/template", map[string]string{"name": name}, &result); err != nil { - return "", err - } - return result.Data, nil -} - -func (c *Client) GetKanikoPod(id uint) (*KanikoPodResponse, error) { - var result Response[KanikoPodResponse] - if err := c.get(ImagesPrefix+"/podname", map[string]string{"id": fmt.Sprintf("%d", id)}, &result); err != nil { - return nil, err - } - return &result.Data, nil -} - -func (c *Client) CreatePipApt(req PipAptBuildRequest) (string, error) { - return c.postString(ImagesPrefix+"/kaniko", req) -} - -func (c *Client) CreateDockerfile(req DockerfileBuildRequest) (string, error) { - return c.postString(ImagesPrefix+"/dockerfile", req) -} - -func (c *Client) CreateEnvd(req EnvdBuildRequest) (string, error) { - return c.postString(ImagesPrefix+"/envd", req) -} - -func (c *Client) RemoveKaniko(ids []uint, admin bool) (string, error) { - path := ImagesPrefix + "/remove" - if admin { - path = AdminImagesPrefix + "/remove" - } - return c.postString(path, IDListRequest{IDList: ids}) -} - -func (c *Client) ListImageRecords(admin bool) (*ListImageResponse, error) { - path := ImagesPrefix + "/image" - if admin { - path = AdminImagesPrefix + "/image" - } - var result Response[ListImageResponse] - if err := c.get(path, nil, &result); err != nil { - return nil, err - } - return &result.Data, nil -} - -func (c *Client) ListAvailableImages() ([]ImageInfo, error) { - var result Response[AvailableImageListResp] - if err := c.get(ImageAvailablePath, nil, &result); err != nil { - return nil, err +func (c *Client) ListImages(available bool) ([]ImageInfo, error) { + if available { + var result Response[AvailableImageListResp] + resp, err := c.httpClient.R(). + SetSuccessResult(&result). + SetErrorResult(&result). + Get(ImageAvailablePath) + if err != nil { + return nil, &NetworkError{Cause: err} + } + if err := errorFromResponse(resp, result.Code, result.Message); err != nil { + return nil, err + } + return result.Data.Images, nil } - return result.Data.Images, nil -} -func (c *Client) UploadImage(req ImageUploadRequest) (string, error) { - return c.postString(ImagesPrefix+"/image", req) -} - -func (c *Client) DeleteImage(id uint) (string, error) { - var result Response[string] - resp, err := c.httpClient.R().SetSuccessResult(&result).SetErrorResult(&result).Delete(fmt.Sprintf("%s/image/%d", ImagesPrefix, id)) + var result Response[ImageListResp] + resp, err := c.httpClient.R(). + SetSuccessResult(&result). + SetErrorResult(&result). + Get(ImageListPath) if err != nil { - return "", &NetworkError{Cause: err} + return nil, &NetworkError{Cause: err} } if err := errorFromResponse(resp, result.Code, result.Message); err != nil { - return "", err - } - return result.Data, nil -} - -func (c *Client) DeleteImages(ids []uint, admin bool) (string, error) { - path := ImagesPrefix + "/deleteimage" - if admin { - path = AdminImagesPrefix + "/deleteimage" - } - return c.postString(path, IDListRequest{IDList: ids}) -} - -func (c *Client) UpdateImageDescription(req ImageDescriptionRequest, admin bool) (string, error) { - return c.postString(imageAdminPath("/description", admin), req) -} - -func (c *Client) UpdateImageType(req ImageTypeRequest, admin bool) (string, error) { - return c.postString(imageAdminPath("/type", admin), req) -} - -func (c *Client) UpdateImageTags(req ImageTagsRequest, admin bool) (string, error) { - return c.postString(imageAdminPath("/tags", admin), req) -} - -func (c *Client) UpdateImageArch(req ImageArchRequest, admin bool) (string, error) { - return c.postString(imageAdminPath("/arch", admin), req) -} - -func (c *Client) TogglePublic(id uint) (string, error) { - return c.postString(fmt.Sprintf("%s/change/%d", AdminImagesPrefix, id), nil) -} - -func (c *Client) ShareImage(req ImageShareRequest) (string, error) { - return c.postString(ImagesPrefix+"/share", req) -} - -func (c *Client) CancelShareImage(req ImageCancelShareRequest) (string, error) { - var result Response[string] - resp, err := c.httpClient.R().SetBody(&req).SetSuccessResult(&result).SetErrorResult(&result).Delete(ImagesPrefix + "/share") - if err != nil { - return "", &NetworkError{Cause: err} - } - if err := errorFromResponse(resp, result.Code, result.Message); err != nil { - return "", err - } - return result.Data, nil -} - -func (c *Client) GetImageGrants(imageID uint) (*ImageGrantResponse, error) { - var result Response[ImageGrantResponse] - if err := c.get(ImagesPrefix+"/share", map[string]string{"imageID": fmt.Sprintf("%d", imageID)}, &result); err != nil { - return nil, err - } - return &result.Data, nil -} - -func (c *Client) ListUngrantedUsers(imageID uint, name string) (*ImageGrantResponse, error) { - var result Response[ImageGrantResponse] - if err := c.get(ImagesPrefix+"/user", map[string]string{"imageID": fmt.Sprintf("%d", imageID), "name": name}, &result); err != nil { - return nil, err - } - return &result.Data, nil -} - -func (c *Client) ListUngrantedAccounts(imageID uint) (*ImageGrantResponse, error) { - var result Response[ImageGrantResponse] - if err := c.get(ImagesPrefix+"/account", map[string]string{"imageID": fmt.Sprintf("%d", imageID)}, &result); err != nil { return nil, err } - return &result.Data, nil -} - -func (c *Client) CheckImageLinks(pairs []ImageInfoLinkPair) (*CheckLinkValidityResponse, error) { - var result Response[CheckLinkValidityResponse] - if err := c.post(ImagesPrefix+"/valid", CheckLinkValidityRequest{LinkPairs: pairs}, &result); err != nil { - return nil, err - } - return &result.Data, nil -} - -func (c *Client) GetHarbor() (*HarborResponse, error) { - var result Response[HarborResponse] - if err := c.get(ImagesPrefix+"/harbor", nil, &result); err != nil { - return nil, err - } - return &result.Data, nil -} - -func (c *Client) GetCredential() (*ProjectCredentialResponse, error) { - var result Response[ProjectCredentialResponse] - if err := c.post(ImagesPrefix+"/credential", nil, &result); err != nil { - return nil, err - } - return &result.Data, nil -} - -func (c *Client) GetQuota() (*ProjectDetailResponse, error) { - var result Response[ProjectDetailResponse] - if err := c.get(ImagesPrefix+"/quota", nil, &result); err != nil { - return nil, err - } - return &result.Data, nil -} - -func (c *Client) UpdateQuota(size int64) (string, error) { - return c.postString(ImagesPrefix+"/quota", map[string]int64{"size": size}) -} - -func (c *Client) ListCudaBaseImages() (*CudaBaseImagesResponse, error) { - var result Response[CudaBaseImagesResponse] - if err := c.get(ImagesPrefix+"/cudabaseimage", nil, &result); err != nil { - return nil, err - } - return &result.Data, nil -} - -func (c *Client) AdminAddCudaBaseImage(req CudaBaseImageRequest) (string, error) { - return c.postString(AdminImagesPrefix+"/cudabaseimage", req) -} - -func (c *Client) AdminDeleteCudaBaseImage(id uint) (string, error) { - var result Response[string] - resp, err := c.httpClient.R().SetSuccessResult(&result).SetErrorResult(&result).Delete(fmt.Sprintf("%s/cudabaseimage/%d", AdminImagesPrefix, id)) - if err != nil { - return "", &NetworkError{Cause: err} - } - if err := errorFromResponse(resp, result.Code, result.Message); err != nil { - return "", err - } - return result.Data, nil -} - -func imageAdminPath(suffix string, admin bool) string { - if admin { - return AdminImagesPrefix + suffix - } - return ImagesPrefix + suffix -} - -func (c *Client) get(path string, params map[string]string, result interface{}) error { - req := c.httpClient.R().SetSuccessResult(result).SetErrorResult(result) - for k, v := range params { - if v != "" { - req.SetQueryParam(k, v) - } - } - resp, err := req.Get(path) - if err != nil { - return &NetworkError{Cause: err} - } - code, msg := responseCodeMsg(result) - return errorFromResponse(resp, code, msg) -} - -func (c *Client) post(path string, body interface{}, result interface{}) error { - resp, err := c.httpClient.R().SetBody(body).SetSuccessResult(result).SetErrorResult(result).Post(path) - if err != nil { - return &NetworkError{Cause: err} - } - code, msg := responseCodeMsg(result) - return errorFromResponse(resp, code, msg) -} - -func (c *Client) postString(path string, body interface{}) (string, error) { - var result Response[string] - if err := c.post(path, body, &result); err != nil { - return "", err - } - return result.Data, nil -} - -func responseCodeMsg(result interface{}) (int, string) { - switch r := result.(type) { - case *Response[ListKanikoResponse]: - return r.Code, r.Message - case *Response[KanikoDetailResponse]: - return r.Code, r.Message - case *Response[KanikoPodResponse]: - return r.Code, r.Message - case *Response[ListImageResponse]: - return r.Code, r.Message - case *Response[AvailableImageListResp]: - return r.Code, r.Message - case *Response[string]: - return r.Code, r.Message - case *Response[ImageGrantResponse]: - return r.Code, r.Message - case *Response[CheckLinkValidityResponse]: - return r.Code, r.Message - case *Response[HarborResponse]: - return r.Code, r.Message - case *Response[ProjectCredentialResponse]: - return r.Code, r.Message - case *Response[ProjectDetailResponse]: - return r.Code, r.Message - case *Response[CudaBaseImagesResponse]: - return r.Code, r.Message - default: - return 0, "" - } + return result.Data.ImageList, nil } diff --git a/cli/internal/api/image_test.go b/cli/internal/api/image_test.go deleted file mode 100644 index e4d8cbdc0..000000000 --- a/cli/internal/api/image_test.go +++ /dev/null @@ -1,150 +0,0 @@ -package api - -import ( - "encoding/json" - "net/http" - "net/http/httptest" - "testing" - - "github.com/imroc/req/v3" -) - -func imageTestClient(t *testing.T, handler http.HandlerFunc) *Client { - t.Helper() - client := NewClient("https://example.invalid") - client.httpClient.GetTransport().WrapRoundTripFunc(func(_ http.RoundTripper) req.HttpRoundTripFunc { - return func(r *http.Request) (*http.Response, error) { - recorder := httptest.NewRecorder() - handler.ServeHTTP(recorder, r) - return recorder.Result(), nil - } - }) - return client -} - -func writeImageTestResponse(t *testing.T, w http.ResponseWriter) { - t.Helper() - w.Header().Set("Content-Type", "application/json") - if err := json.NewEncoder(w).Encode(map[string]interface{}{ - "code": 0, - "data": nil, - "msg": "", - }); err != nil { - t.Fatalf("encode response: %v", err) - } -} - -func TestImageClientRoutesMatchBackend(t *testing.T) { - tests := []struct { - name string - method string - path string - query string - call func(*Client) error - }{ - {"list user builds", http.MethodGet, "/api/v1/images/kaniko", "", func(c *Client) error { _, err := c.ListKaniko(false); return err }}, - {"list admin builds", http.MethodGet, "/api/v1/admin/images/kaniko", "", func(c *Client) error { _, err := c.ListKaniko(true); return err }}, - {"get build", http.MethodGet, "/api/v1/images/getbyname", "name=build-1", func(c *Client) error { _, err := c.GetKanikoByName("build-1"); return err }}, - {"get build template", http.MethodGet, "/api/v1/images/template", "name=build-1", func(c *Client) error { _, err := c.GetKanikoTemplateByName("build-1"); return err }}, - {"get build pod", http.MethodGet, "/api/v1/images/podname", "id=7", func(c *Client) error { _, err := c.GetKanikoPod(7); return err }}, - {"create pip apt build", http.MethodPost, "/api/v1/images/kaniko", "", func(c *Client) error { _, err := c.CreatePipApt(PipAptBuildRequest{}); return err }}, - {"create dockerfile build", http.MethodPost, "/api/v1/images/dockerfile", "", func(c *Client) error { _, err := c.CreateDockerfile(DockerfileBuildRequest{}); return err }}, - {"create envd build", http.MethodPost, "/api/v1/images/envd", "", func(c *Client) error { _, err := c.CreateEnvd(EnvdBuildRequest{}); return err }}, - {"remove user builds", http.MethodPost, "/api/v1/images/remove", "", func(c *Client) error { _, err := c.RemoveKaniko([]uint{7}, false); return err }}, - {"remove admin builds", http.MethodPost, "/api/v1/admin/images/remove", "", func(c *Client) error { _, err := c.RemoveKaniko([]uint{7}, true); return err }}, - {"list user images", http.MethodGet, "/api/v1/images/image", "", func(c *Client) error { _, err := c.ListImageRecords(false); return err }}, - {"list admin images", http.MethodGet, "/api/v1/admin/images/image", "", func(c *Client) error { _, err := c.ListImageRecords(true); return err }}, - {"list available images", http.MethodGet, "/api/v1/images/available", "", func(c *Client) error { _, err := c.ListAvailableImages(); return err }}, - {"upload image", http.MethodPost, "/api/v1/images/image", "", func(c *Client) error { _, err := c.UploadImage(ImageUploadRequest{}); return err }}, - {"delete image", http.MethodDelete, "/api/v1/images/image/7", "", func(c *Client) error { _, err := c.DeleteImage(7); return err }}, - {"delete user images", http.MethodPost, "/api/v1/images/deleteimage", "", func(c *Client) error { _, err := c.DeleteImages([]uint{7}, false); return err }}, - {"delete admin images", http.MethodPost, "/api/v1/admin/images/deleteimage", "", func(c *Client) error { _, err := c.DeleteImages([]uint{7}, true); return err }}, - {"update user description", http.MethodPost, "/api/v1/images/description", "", func(c *Client) error { - _, err := c.UpdateImageDescription(ImageDescriptionRequest{}, false) - return err - }}, - {"update admin description", http.MethodPost, "/api/v1/admin/images/description", "", func(c *Client) error { _, err := c.UpdateImageDescription(ImageDescriptionRequest{}, true); return err }}, - {"update user type", http.MethodPost, "/api/v1/images/type", "", func(c *Client) error { _, err := c.UpdateImageType(ImageTypeRequest{}, false); return err }}, - {"update admin type", http.MethodPost, "/api/v1/admin/images/type", "", func(c *Client) error { _, err := c.UpdateImageType(ImageTypeRequest{}, true); return err }}, - {"update user tags", http.MethodPost, "/api/v1/images/tags", "", func(c *Client) error { _, err := c.UpdateImageTags(ImageTagsRequest{}, false); return err }}, - {"update admin tags", http.MethodPost, "/api/v1/admin/images/tags", "", func(c *Client) error { _, err := c.UpdateImageTags(ImageTagsRequest{}, true); return err }}, - {"update user arch", http.MethodPost, "/api/v1/images/arch", "", func(c *Client) error { _, err := c.UpdateImageArch(ImageArchRequest{}, false); return err }}, - {"update admin arch", http.MethodPost, "/api/v1/admin/images/arch", "", func(c *Client) error { _, err := c.UpdateImageArch(ImageArchRequest{}, true); return err }}, - {"toggle public", http.MethodPost, "/api/v1/admin/images/change/7", "", func(c *Client) error { _, err := c.TogglePublic(7); return err }}, - {"share image", http.MethodPost, "/api/v1/images/share", "", func(c *Client) error { _, err := c.ShareImage(ImageShareRequest{}); return err }}, - {"cancel share", http.MethodDelete, "/api/v1/images/share", "", func(c *Client) error { _, err := c.CancelShareImage(ImageCancelShareRequest{}); return err }}, - {"list grants", http.MethodGet, "/api/v1/images/share", "imageID=7", func(c *Client) error { _, err := c.GetImageGrants(7); return err }}, - {"list ungranted users", http.MethodGet, "/api/v1/images/user", "imageID=7&name=alice", func(c *Client) error { _, err := c.ListUngrantedUsers(7, "alice"); return err }}, - {"list ungranted accounts", http.MethodGet, "/api/v1/images/account", "imageID=7", func(c *Client) error { _, err := c.ListUngrantedAccounts(7); return err }}, - {"validate links", http.MethodPost, "/api/v1/images/valid", "", func(c *Client) error { _, err := c.CheckImageLinks(nil); return err }}, - {"get harbor", http.MethodGet, "/api/v1/images/harbor", "", func(c *Client) error { _, err := c.GetHarbor(); return err }}, - {"get credential", http.MethodPost, "/api/v1/images/credential", "", func(c *Client) error { _, err := c.GetCredential(); return err }}, - {"get quota", http.MethodGet, "/api/v1/images/quota", "", func(c *Client) error { _, err := c.GetQuota(); return err }}, - {"update quota", http.MethodPost, "/api/v1/images/quota", "", func(c *Client) error { _, err := c.UpdateQuota(1024); return err }}, - {"list cuda images", http.MethodGet, "/api/v1/images/cudabaseimage", "", func(c *Client) error { _, err := c.ListCudaBaseImages(); return err }}, - {"add cuda image", http.MethodPost, "/api/v1/admin/images/cudabaseimage", "", func(c *Client) error { _, err := c.AdminAddCudaBaseImage(CudaBaseImageRequest{}); return err }}, - {"delete cuda image", http.MethodDelete, "/api/v1/admin/images/cudabaseimage/7", "", func(c *Client) error { _, err := c.AdminDeleteCudaBaseImage(7); return err }}, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - client := imageTestClient(t, func(w http.ResponseWriter, r *http.Request) { - if r.Method != tt.method || r.URL.Path != tt.path { - t.Errorf("request = %s %s, want %s %s", r.Method, r.URL.Path, tt.method, tt.path) - } - if r.URL.RawQuery != tt.query { - t.Errorf("query = %q, want %q", r.URL.RawQuery, tt.query) - } - writeImageTestResponse(t, w) - }) - if err := tt.call(client); err != nil { - t.Fatalf("call failed: %v", err) - } - }) - } -} - -func TestImageClientShareBodiesMatchBackendDTOs(t *testing.T) { - tests := []struct { - name string - call func(*Client) error - want map[string]interface{} - }{ - { - name: "share", - call: func(c *Client) error { - _, err := c.ShareImage(ImageShareRequest{IDList: []uint{2, 3}, ImageID: 7, Type: "account"}) - return err - }, - want: map[string]interface{}{"idList": []interface{}{float64(2), float64(3)}, "imageID": float64(7), "type": "account"}, - }, - { - name: "cancel", - call: func(c *Client) error { - _, err := c.CancelShareImage(ImageCancelShareRequest{ID: 2, ImageID: 7, Type: "user"}) - return err - }, - want: map[string]interface{}{"id": float64(2), "imageID": float64(7), "type": "user"}, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - client := imageTestClient(t, func(w http.ResponseWriter, r *http.Request) { - var got map[string]interface{} - if err := json.NewDecoder(r.Body).Decode(&got); err != nil { - t.Fatalf("decode body: %v", err) - } - gotJSON, _ := json.Marshal(got) - wantJSON, _ := json.Marshal(tt.want) - if string(gotJSON) != string(wantJSON) { - t.Errorf("body = %s, want %s", gotJSON, wantJSON) - } - writeImageTestResponse(t, w) - }) - if err := tt.call(client); err != nil { - t.Fatalf("call failed: %v", err) - } - }) - } -} diff --git a/cli/internal/i18n/catalog_image.go b/cli/internal/i18n/catalog_image.go deleted file mode 100644 index 5dc4f5359..000000000 --- a/cli/internal/i18n/catalog_image.go +++ /dev/null @@ -1,182 +0,0 @@ -package i18n - -var catalogImage = map[Language]map[string]string{ - En: { - "image_build_short": "Manage image builds", - "image_build_ls_short": "List image build records", - "image_build_pip-apt_short": "Build an image from base image plus pip/apt packages", - "image_build_dockerfile_short": "Build an image from Dockerfile", - "image_build_envd_short": "Build an image from envd", - "image_build_remove_short": "Cancel or remove image build records", - "image_build_get_short": "Get an image build record", - "image_build_template_short": "Get image build template", - "image_build_pod_short": "Get image build pod", - "image_upload_short": "Upload/register an existing image link", - "image_delete_short": "Delete an image", - "image_delete-many_short": "Delete multiple images", - "image_description_short": "Update image description", - "image_type_short": "Update image task type", - "image_tags_short": "Update image tags", - "image_arch_short": "Update image architectures", - "image_valid_short": "Validate image links", - "image_share_short": "Manage image sharing", - "image_share_ls_short": "List image grants", - "image_share_users_short": "List users not granted an image", - "image_share_accounts_short": "List accounts not granted an image", - "image_share_add_short": "Share an image with users or accounts", - "image_share_remove_short": "Cancel image sharing", - "image_cuda_short": "View CUDA base images", - "image_cuda_ls_short": "List CUDA base images", - "image_harbor_short": "View Harbor information", - "image_harbor_info_short": "Get Harbor address", - "image_harbor_credential_short": "Create and show Harbor project credentials", - "image_quota_short": "View or update Harbor project quota", - "image_quota_get_short": "Get Harbor project quota", - "image_quota_set_short": "Update Harbor project quota", - "admin_image_short": "Manage admin image resources", - "admin_image_build-ls_short": "List all image build records", - "admin_image_build-remove_short": "Cancel or remove image build records", - "admin_image_ls_short": "List all images", - "admin_image_delete-many_short": "Delete multiple images", - "admin_image_description_short": "Update image description", - "admin_image_type_short": "Update image task type", - "admin_image_tags_short": "Update image tags", - "admin_image_arch_short": "Update image architectures", - "admin_image_public_short": "Toggle public visibility", - "admin_image_cuda_short": "Manage CUDA base images", - "admin_image_cuda_add_short": "Add a CUDA base image", - "admin_image_cuda_delete_short": "Delete a CUDA base image", - - "image_flag_name": "Image name", - "image_flag_tag": "Image tag", - "image_flag_image": "Base image or image link", - "image_flag_description": "Description", - "image_flag_packages": "APT packages, separated by spaces", - "image_flag_requirements": "Python requirements content", - "image_flag_file": "Read content from file", - "image_flag_dockerfile": "Dockerfile content", - "image_flag_envd": "envd content", - "image_flag_python": "Python version", - "image_flag_base": "envd base image", - "image_flag_build-source": "envd build source", - "image_flag_tags": "Comma-separated tags", - "image_flag_archs": "Comma-separated architectures", - "image_flag_type": "Image task type", - "image_flag_ids": "Comma-separated IDs", - "image_flag_share-type": "Share type: user or account", - "image_flag_target-id": "Share target ID", - "image_flag_size": "Quota size", - "image_flag_label": "Display label", - "image_flag_image-label": "Image label", - "image_flag_value": "Image value", - "image_flag_template": "Template text", - "image_flag_links": "Comma-separated image links", - "image_label_id": "image ID", - "image_label_name": "image name", - "err_invalid_ids": "invalid ID list: %s", - "err_confirm_needed": "confirmation is required (--yes)", - "err_invalid_image_arch": "invalid image architecture: %s", - "err_invalid_image_value": "invalid %s: %s", - "err_mutually_exclusive_flags": "--%s and --%s cannot be used together", - "image_success": "Image operation completed: %s", - "image_output_harbor": "Harbor", - "image_output_username": "Username", - "image_output_password": "Password", - "image_output_project": "Project", - "image_output_used": "Used", - "image_output_quota": "Quota", - "table_label": "LABEL", - "table_image_label": "IMAGE_LABEL", - "table_value": "VALUE", - "table_nickname": "NICKNAME", - }, - ZhCN: { - "image_build_short": "管理镜像构建", - "image_build_ls_short": "列出镜像构建记录", - "image_build_pip-apt_short": "基于基础镜像和 pip/apt 包构建镜像", - "image_build_dockerfile_short": "基于 Dockerfile 构建镜像", - "image_build_envd_short": "基于 envd 构建镜像", - "image_build_remove_short": "取消或删除镜像构建记录", - "image_build_get_short": "查看镜像构建记录", - "image_build_template_short": "查看镜像构建模板", - "image_build_pod_short": "查看镜像构建 Pod", - "image_upload_short": "上传/登记已有镜像链接", - "image_delete_short": "删除镜像", - "image_delete-many_short": "批量删除镜像", - "image_description_short": "更新镜像描述", - "image_type_short": "更新镜像任务类型", - "image_tags_short": "更新镜像标签", - "image_arch_short": "更新镜像架构", - "image_valid_short": "校验镜像链接", - "image_share_short": "管理镜像分享", - "image_share_ls_short": "列出镜像分享对象", - "image_share_users_short": "列出未授权用户", - "image_share_accounts_short": "列出未授权账户", - "image_share_add_short": "分享镜像给用户或账户", - "image_share_remove_short": "取消镜像分享", - "image_cuda_short": "查看 CUDA 基础镜像", - "image_cuda_ls_short": "列出 CUDA 基础镜像", - "image_harbor_short": "查看 Harbor 信息", - "image_harbor_info_short": "查看 Harbor 地址", - "image_harbor_credential_short": "创建并显示 Harbor 项目凭据", - "image_quota_short": "查看或更新 Harbor 项目配额", - "image_quota_get_short": "查看 Harbor 项目配额", - "image_quota_set_short": "更新 Harbor 项目配额", - "admin_image_short": "管理管理员镜像资源", - "admin_image_build-ls_short": "列出所有镜像构建记录", - "admin_image_build-remove_short": "取消或删除镜像构建记录", - "admin_image_ls_short": "列出所有镜像", - "admin_image_delete-many_short": "批量删除镜像", - "admin_image_description_short": "更新镜像描述", - "admin_image_type_short": "更新镜像任务类型", - "admin_image_tags_short": "更新镜像标签", - "admin_image_arch_short": "更新镜像架构", - "admin_image_public_short": "切换公共可见性", - "admin_image_cuda_short": "管理 CUDA 基础镜像", - "admin_image_cuda_add_short": "添加 CUDA 基础镜像", - "admin_image_cuda_delete_short": "删除 CUDA 基础镜像", - - "image_flag_name": "镜像名称", - "image_flag_tag": "镜像标签", - "image_flag_image": "基础镜像或镜像链接", - "image_flag_description": "描述", - "image_flag_packages": "APT 包，使用空格分隔", - "image_flag_requirements": "Python requirements 内容", - "image_flag_file": "从文件读取内容", - "image_flag_dockerfile": "Dockerfile 内容", - "image_flag_envd": "envd 内容", - "image_flag_python": "Python 版本", - "image_flag_base": "envd 基础镜像", - "image_flag_build-source": "envd 构建来源", - "image_flag_tags": "逗号分隔标签", - "image_flag_archs": "逗号分隔架构", - "image_flag_type": "镜像任务类型", - "image_flag_ids": "逗号分隔 ID", - "image_flag_share-type": "分享类型：user 或 account", - "image_flag_target-id": "分享目标 ID", - "image_flag_size": "配额大小", - "image_flag_label": "展示标签", - "image_flag_image-label": "镜像标签", - "image_flag_value": "镜像值", - "image_flag_template": "模板文本", - "image_flag_links": "逗号分隔镜像链接", - "image_label_id": "镜像 ID", - "image_label_name": "镜像名称", - "err_invalid_ids": "无效的 ID 列表：%s", - "err_confirm_needed": "必须确认操作 (--yes)", - "err_invalid_image_arch": "无效的镜像架构：%s", - "err_invalid_image_value": "无效的 %s：%s", - "err_mutually_exclusive_flags": "--%s 和 --%s 不能同时使用", - "image_success": "镜像操作完成：%s", - "image_output_harbor": "Harbor", - "image_output_username": "用户名", - "image_output_password": "密码", - "image_output_project": "项目", - "image_output_used": "已用", - "image_output_quota": "配额", - "table_label": "标签", - "table_image_label": "镜像标签", - "table_value": "值", - "table_nickname": "昵称", - }, -} diff --git a/cli/internal/i18n/catalog_read.go b/cli/internal/i18n/catalog_read.go index 74bd657f2..7db3feb20 100644 --- a/cli/internal/i18n/catalog_read.go +++ b/cli/internal/i18n/catalog_read.go @@ -17,8 +17,8 @@ var catalogRead = map[Language]map[string]string{ "job_events_short": "List events for a job", "job_yaml_short": "Show job YAML", - "image_short": "Manage images and image builds", - "image_long": "Build, upload, delete, share, and update Crater images and image build records.", + "image_short": "View images", + "image_long": "View container image lists from the active Crater platform.", "image_ls_short": "List images", "account_short": "View accounts", @@ -182,8 +182,8 @@ var catalogRead = map[Language]map[string]string{ "job_events_short": "列出作业事件", "job_yaml_short": "显示作业 YAML", - "image_short": "管理镜像与镜像构建", - "image_long": "构建、上传、删除、分享和更新 Crater 镜像与镜像构建记录。", + "image_short": "查看镜像", + "image_long": "从当前激活的 Crater 平台查看容器镜像列表。", "image_ls_short": "列出镜像", "account_short": "查看账户", diff --git a/cli/internal/i18n/i18n.go b/cli/internal/i18n/i18n.go index 264a91e85..12c35f242 100644 --- a/cli/internal/i18n/i18n.go +++ b/cli/internal/i18n/i18n.go @@ -24,7 +24,6 @@ var translations = mergeCatalogs( catalogCompletion, catalogDownload, catalogRead, - catalogImage, catalogOrder, catalogErrors, catalogJob, diff --git a/cli/skills/crater-cli-admin-image-management/SKILL.md b/cli/skills/crater-cli-admin-image-management/SKILL.md deleted file mode 100644 index d69cade0f..000000000 --- a/cli/skills/crater-cli-admin-image-management/SKILL.md +++ /dev/null @@ -1,38 +0,0 @@ ---- -name: crater-cli-admin-image-management -version: 1.0.0 -description: "Crater CLI 管理员镜像管理：指导 AI Agent 使用 crater admin image 查看、删除、修改所有用户镜像、镜像构建记录与 CUDA base image。仅在用户明确要求管理员/平台级镜像操作时使用。" -metadata: - requires: - bins: ["crater"] - cliHelp: "crater admin image --help" ---- - -# Crater CLI 管理员镜像管理 - -**CRITICAL — 开始前 MUST 先读取 `crater-cli-shared`（可能路径：[`../crater-cli-shared/SKILL.md`](../crater-cli-shared/SKILL.md)）。** - -本 Skill 仅用于管理员镜像操作。普通用户镜像构建、上传、分享和 Harbor 凭据流程请使用 `crater-cli-image-management`。 - -## 命令 - -```bash -crater admin image build-ls --json -crater admin image build-remove --ids 1,2 --json --no-interactive -crater admin image ls --json -crater admin image delete-many --ids 1,2 --json --no-interactive -crater admin image description 1 --description "Updated" --json --no-interactive -crater admin image type 1 --type jupyter --json --no-interactive -crater admin image tags 1 --tags CUDA,Jupyter --json --no-interactive -crater admin image arch 1 --archs linux/amd64 --json --no-interactive -crater admin image public 1 --json --no-interactive -crater admin image cuda add --image-label cuda124 --label "CUDA 12.4" --value registry/nvidia/cuda:12.4 --json --no-interactive -crater admin image cuda delete 1 --json --no-interactive -``` - -## 规则 - -- 管理员命令统一使用 `crater admin image ...` 前缀；不要使用 `--admin`。 -- 修改/删除操作会影响平台级镜像资源，执行前确认用户明确要求管理员操作。 -- 优先使用 `--json --no-interactive`，读取 `stdout.data.message` 或对应数据键。 -- 403 表示 active credentials 不是平台管理员或权限不足。 diff --git a/cli/skills/crater-cli-image-management/SKILL.md b/cli/skills/crater-cli-image-management/SKILL.md deleted file mode 100644 index 29286bb98..000000000 --- a/cli/skills/crater-cli-image-management/SKILL.md +++ /dev/null @@ -1,73 +0,0 @@ ---- -name: crater-cli-image-management -version: 1.0.0 -description: "Crater CLI 用户镜像与环境管理：指导 AI Agent 使用 crater image 构建、上传、删除、分享、更新自己可管理的镜像、查看 CUDA base image 和获取 Harbor 凭据。管理员镜像操作请使用 crater-cli-admin-image-management。" -metadata: - requires: - bins: ["crater"] - cliHelp: "crater image --help" ---- - -# Crater CLI 镜像与环境管理 - -**CRITICAL — 开始前 MUST 先读取 `crater-cli-shared`（可能路径：[`../crater-cli-shared/SKILL.md`](../crater-cli-shared/SKILL.md)），其中包含全局选项、非交互调用、错误处理和敏感信息规则。** - -通过 `crater image` 管理当前用户可见/可操作的 Crater 镜像、镜像构建和 Harbor 项目时，遵守本规则。 - -## 适用场景 - -- 用户需要用 pip/apt、Dockerfile 或 envd 构建镜像。 -- 用户需要上传/登记已有镜像链接。 -- 用户需要删除镜像或取消/删除镜像构建任务。 -- 用户需要修改镜像描述、类型、标签或架构。 -- 用户需要分享或取消分享镜像。 -- 用户需要查看 CUDA base image。 -- 用户需要查看 Harbor 地址、配额或生成 Harbor 项目凭据。 - -## 安全边界 - -- 本领域大部分命令会修改平台状态。 -- 需要脚本化时优先使用 `--json --no-interactive`。 -- `crater image harbor credential` 会创建并输出 Harbor 凭据，必须显式加 `--yes`，不要在不安全日志中暴露输出。 -- 不要在用户侧命令使用 `--admin`。管理员镜像操作统一使用 `crater admin image ...`，并切换到 `crater-cli-admin-image-management`。 - -## 常用范例 - -```bash -crater image build pip-apt \ - --name cuda-demo \ - --tag v1 \ - --image nvidia/cuda:12.4.1-devel-ubuntu22.04 \ - --packages "git vim" \ - --requirements "torch==2.4.0" \ - --tags CUDA,Pytorch \ - --json --no-interactive - -crater image build dockerfile --name custom --tag v1 --file ./Dockerfile --json --no-interactive -crater image build envd --name envd-demo --tag v1 --file ./build.envd --json --no-interactive -crater image upload --image registry/project/repo:tag --type custom --json --no-interactive -crater image delete-many --ids 1,2 --json --no-interactive -crater image description 1 --description "Updated description" --json --no-interactive -crater image type 1 --type jupyter --json --no-interactive -crater image tags 1 --tags CUDA,Jupyter --json --no-interactive -crater image arch 1 --archs linux/amd64 --json --no-interactive -crater image share add 1 --share-type user --ids 10,11 --json --no-interactive -crater image share remove 1 --share-type user --target-id 10 --json --no-interactive -crater image cuda ls --json --no-interactive -crater image harbor credential --yes --json --no-interactive -``` - -## 工作流 - -1. 先用 `crater auth ls --json` 确认 active credentials。 -2. 构建镜像前确认 `--name`、`--tag` 和构建内容非空。 -3. Dockerfile/envd 大内容优先通过 `--file` 传入，避免 shell 转义问题。 -4. 批量删除或取消构建时使用逗号分隔 ID。 -5. 处理 Harbor 凭据时只在用户明确要求时执行，并提醒其输出敏感。 -6. 如果用户要求管理 CUDA base image、平台级镜像列表、删除他人镜像或修改公共可见性，切换到管理员 skill。 - -## 输出处理 - -- 成功 JSON 读取 `stdout.data.message` 或具体数据键。 -- API 失败读取 stderr JSON 的 `category`、`code`、`context.http_status`。 -- 401 表示需要重新登录；403 通常表示权限不足。 diff --git a/cli/test/snapshots/image/image_test.go b/cli/test/snapshots/image/image_test.go deleted file mode 100644 index 9d36f8b10..000000000 --- a/cli/test/snapshots/image/image_test.go +++ /dev/null @@ -1,87 +0,0 @@ -package image_test - -import ( - "os" - "strings" - "testing" - - "github.com/raids-lab/crater/cli/internal/snaptest" -) - -const goldenStemImage = "image" - -func TestImageSnapshotsEN(t *testing.T) { - runImageSnapshots(t, "en") -} - -func TestImageSnapshotsZhCN(t *testing.T) { - runImageSnapshots(t, "zh-CN") -} - -func TestImageUnknownSubcommandText(t *testing.T) { - home := t.TempDir() - result, err := snaptest.Run(snaptest.CraterExecutable(t), snaptest.EnvMinimal(home, "en"), []string{"image", "list", "--no-interactive"}) - if err != nil { - t.Fatalf("run image list: %v", err) - } - if result.ExitCode != 2 { - t.Fatalf("exit code = %d, want 2", result.ExitCode) - } - if !strings.Contains(result.Stderr, `unknown command "list" for "crater image"`) { - t.Fatalf("stderr missing unknown-command error: %q", result.Stderr) - } -} - -func runImageSnapshots(t *testing.T, lang string) { - t.Helper() - path := snaptest.GoldenFileT(t, "image", goldenStemImage, lang) - home := t.TempDir() - baseEnv := snaptest.EnvMinimal(home, lang) - bin := snaptest.CraterExecutable(t) - cases := []snaptest.Case{ - {ID: "01-image-typo-json", Args: []string{"image", "list", "--json", "--no-interactive"}}, - {ID: "02-build-pip-apt-missing-json", Args: []string{"image", "build", "pip-apt", "--json", "--no-interactive"}}, - {ID: "03-build-dockerfile-missing-content-json", Args: []string{"image", "build", "dockerfile", "--name", "img", "--tag", "v1", "--json", "--no-interactive"}}, - {ID: "04-build-envd-invalid-source-json", Args: []string{"image", "build", "envd", "--name", "img", "--tag", "v1", "--envd", "x", "--build-source", "bad", "--json", "--no-interactive"}}, - {ID: "05-upload-missing-image-json", Args: []string{"image", "upload", "--json", "--no-interactive"}}, - {ID: "06-delete-invalid-id-json", Args: []string{"image", "delete", "abc", "--json", "--no-interactive"}}, - {ID: "07-delete-many-invalid-ids-json", Args: []string{"image", "delete-many", "--ids", "1,x", "--json", "--no-interactive"}}, - {ID: "08-share-add-missing-ids-json", Args: []string{"image", "share", "add", "1", "--json", "--no-interactive"}}, - {ID: "09-share-remove-missing-target-json", Args: []string{"image", "share", "remove", "1", "--json", "--no-interactive"}}, - {ID: "10-harbor-credential-confirm-json", Args: []string{"image", "harbor", "credential", "--json", "--no-interactive"}}, - {ID: "11-cuda-add-missing-json", Args: []string{"image", "cuda", "add", "--json", "--no-interactive"}}, - {ID: "12-valid-missing-links-json", Args: []string{"image", "valid", "--json", "--no-interactive"}}, - {ID: "13-build-get-missing-json", Args: []string{"image", "build", "get", "--json", "--no-interactive"}}, - {ID: "14-build-pod-invalid-json", Args: []string{"image", "build", "pod", "bad", "--json", "--no-interactive"}}, - {ID: "15-admin-build-remove-missing-json", Args: []string{"admin", "image", "build-remove", "--json", "--no-interactive"}}, - {ID: "16-admin-type-invalid-json", Args: []string{"admin", "image", "type", "1", "--type", "all", "--json", "--no-interactive"}}, - {ID: "17-image-ls-404-json", Args: []string{"image", "ls", "--json", "--no-interactive"}}, - {ID: "18-image-available-404-json", Args: []string{"image", "ls", "--available", "--json", "--no-interactive"}}, - {ID: "19-build-pip-apt-404-json", Args: []string{"image", "build", "pip-apt", "--name", "img", "--tag", "v1", "--image", "base:latest", "--json", "--no-interactive"}}, - {ID: "20-share-ls-404-json", Args: []string{"image", "share", "ls", "1", "--json", "--no-interactive"}}, - {ID: "21-admin-image-ls-404-json", Args: []string{"admin", "image", "ls", "--json", "--no-interactive"}}, - {ID: "22-build-dockerfile-mutually-exclusive-json", Args: []string{"image", "build", "dockerfile", "--name", "img", "--tag", "v1", "--dockerfile", "FROM scratch", "--file", "Dockerfile", "--json", "--no-interactive"}}, - {ID: "23-image-arch-missing-json", Args: []string{"image", "arch", "1", "--json", "--no-interactive"}}, - {ID: "24-admin-cuda-add-missing-json", Args: []string{"admin", "image", "cuda", "add", "--json", "--no-interactive"}}, - {ID: "25-image-help", Args: []string{"image", "--help"}}, - {ID: "26-image-upload-help", Args: []string{"image", "upload", "--help"}}, - {ID: "27-image-build-pip-apt-help", Args: []string{"image", "build", "pip-apt", "--help"}}, - } - results := make([]*snaptest.Result, len(cases)) - for i := range cases { - env := baseEnv - switch cases[i].ID { - case "17-image-ls-404-json", "18-image-available-404-json", "19-build-pip-apt-404-json", "20-share-ls-404-json", "21-admin-image-ls-404-json": - env = append(baseEnv, "CRATER_TEST_SANDBOX_HTTP=error404") - } - r, err := snaptest.Run(bin, env, cases[i].Args) - if err != nil { - t.Fatalf("case %s: %v", cases[i].ID, err) - } - results[i] = r - } - update := os.Getenv("UPDATE_SNAPSHOTS") == "1" || os.Getenv("UPDATE_SNAPSHOTS") == "true" - if err := snaptest.MatchOrUpdateGolden(path, lang, cases, results, update); err != nil { - t.Fatal(err) - } -} diff --git a/cli/test/snapshots/read/read_matrix_test.go b/cli/test/snapshots/read/read_matrix_test.go index 765ae1cf3..399712982 100644 --- a/cli/test/snapshots/read/read_matrix_test.go +++ b/cli/test/snapshots/read/read_matrix_test.go @@ -16,10 +16,7 @@ func TestReadCommandMatrix(t *testing.T) { {"node", "ls"}, {"node", "get"}, {"node", "pods"}, {"node", "gpu"}, {"job", "ls"}, {"job", "get"}, {"job", "pods"}, {"job", "events"}, {"job", "yaml"}, {"job", "template"}, {"job", "token"}, {"job", "secret"}, {"job", "ssh"}, {"job", "snapshot"}, {"job", "alert"}, {"job", "delete"}, {"job", "create", "jupyter"}, {"job", "create", "webide"}, {"job", "create", "custom"}, {"job", "create", "tensorflow"}, {"job", "create", "pytorch"}, - {"image", "ls"}, {"image", "build", "ls"}, {"image", "build", "pip-apt"}, {"image", "build", "dockerfile"}, {"image", "build", "envd"}, {"image", "build", "remove"}, {"image", "build", "get"}, {"image", "build", "template"}, {"image", "build", "pod"}, - {"image", "upload"}, {"image", "delete"}, {"image", "delete-many"}, {"image", "description"}, {"image", "type"}, {"image", "tags"}, {"image", "arch"}, {"image", "valid"}, - {"image", "share", "ls"}, {"image", "share", "users"}, {"image", "share", "accounts"}, {"image", "share", "add"}, {"image", "share", "remove"}, - {"image", "cuda", "ls"}, {"image", "harbor", "info"}, {"image", "harbor", "credential"}, {"image", "quota", "get"}, {"image", "quota", "set"}, + {"image", "ls"}, {"account", "ls"}, {"account", "get"}, {"account", "members"}, {"account", "users-out"}, {"account", "billing", "config"}, {"account", "billing", "members"}, {"resource", "ls"}, {"resource", "networks"}, {"resource", "vgpu"}, {"resource", "prices"}, {"dataset", "ls"}, {"dataset", "get"}, {"dataset", "users"}, {"dataset", "queues"}, {"dataset", "users-out"}, {"dataset", "queues-out"}, @@ -37,7 +34,6 @@ func TestReadCommandMatrix(t *testing.T) { {"admin", "dataset", "ls"}, {"admin", "model-download", "ls"}, {"admin", "billing", "status"}, {"admin", "billing", "jobs"}, - {"admin", "image", "build-ls"}, {"admin", "image", "build-remove"}, {"admin", "image", "ls"}, {"admin", "image", "delete-many"}, {"admin", "image", "description"}, {"admin", "image", "type"}, {"admin", "image", "tags"}, {"admin", "image", "arch"}, {"admin", "image", "public"}, {"admin", "image", "cuda", "add"}, {"admin", "image", "cuda", "delete"}, {"admin", "job", "ls"}, {"admin", "job", "delete"}, {"admin", "job", "lock"}, {"admin", "job", "unlock"}, {"admin", "job", "keep"}, {"admin", "job", "clean", "waiting-jupyter"}, {"admin", "job", "clean", "waiting-custom"}, {"admin", "job", "clean", "long-running"}, {"admin", "job", "clean", "low-gpu"}, {"admin", "order", "ls"}, {"admin", "order", "get"}, {"admin", "order", "approve"}, {"admin", "order", "reject"}, {"admin", "order", "check"}, @@ -56,7 +52,7 @@ func TestReadCommandMatrix(t *testing.T) { } apiCases := [][]string{ - {"node", "ls"}, {"job", "ls"}, {"image", "ls"}, {"image", "build", "ls"}, {"image", "cuda", "ls"}, {"image", "harbor", "info"}, {"image", "quota", "get"}, + {"node", "ls"}, {"job", "ls"}, {"image", "ls"}, {"account", "ls"}, {"resource", "ls"}, {"dataset", "ls"}, {"template", "ls"}, {"model-download", "ls"}, {"context", "prequeue"}, {"context", "quota"}, {"context", "resources"}, {"context", "billing"}, {"billing", "status"}, {"billing", "summary"}, {"billing", "prices"}, {"billing", "jobs"}, @@ -64,7 +60,7 @@ func TestReadCommandMatrix(t *testing.T) { {"admin", "system-config", "llm"}, {"admin", "system-config", "gpu-analysis"}, {"admin", "system-config", "prequeue"}, {"admin", "queue-quotas"}, {"admin", "gpu-analyses"}, {"admin", "operation-logs"}, {"admin", "cronjobs"}, {"admin", "whitelist"}, {"admin", "account", "ls"}, {"admin", "dataset", "ls"}, {"admin", "model-download", "ls"}, - {"admin", "billing", "status"}, {"admin", "billing", "jobs"}, {"admin", "image", "ls"}, {"admin", "image", "build-ls"}, {"admin", "job", "ls"}, {"admin", "order", "ls"}, {"admin", "user", "ls"}, {"admin", "user", "billing", "summary"}, + {"admin", "billing", "status"}, {"admin", "billing", "jobs"}, {"admin", "job", "ls"}, {"admin", "order", "ls"}, {"admin", "user", "ls"}, {"admin", "user", "billing", "summary"}, } env404 := append(baseEnv, "CRATER_TEST_SANDBOX_HTTP=error404") for _, command := range apiCases { diff --git a/cli/testdata/snapshots/image/image.en.txtar b/cli/testdata/snapshots/image/image.en.txtar deleted file mode 100644 index c972e0ad1..000000000 --- a/cli/testdata/snapshots/image/image.en.txtar +++ /dev/null @@ -1,405 +0,0 @@ -# Crater CLI snapshot bundle (txtar). Regenerate: make snapshot-update (or UPDATE_SNAPSHOTS=1 go test ./test/snapshots/...) --- en/01-image-typo-json/argv -- -crater image list --json --no-interactive --- en/01-image-typo-json/exit -- -2 --- en/01-image-typo-json/stdout -- --- en/01-image-typo-json/stderr -- -{ - "category": "usage_error", - "code": "ERR_UNKNOWN_COMMAND", - "message": "unknown command \"list\" for \"crater image\"\n\nDid you mean this?\n\tls\n\nRun \"crater image --help\" for usage." -} --- en/02-build-pip-apt-missing-json/argv -- -crater image build pip-apt --json --no-interactive --- en/02-build-pip-apt-missing-json/exit -- -2 --- en/02-build-pip-apt-missing-json/stdout -- --- en/02-build-pip-apt-missing-json/stderr -- -{ - "category": "usage_error", - "code": "ERR_MISSING_REQUIRED_FLAG", - "message": "Image name is required (--name)\nImage tag is required (--tag)", - "context": { - "issues": [ - { - "code": "ERR_MISSING_REQUIRED_FLAG", - "field": "name", - "message": "Image name is required (--name)" - }, - { - "code": "ERR_MISSING_REQUIRED_FLAG", - "field": "tag", - "message": "Image tag is required (--tag)" - } - ] - } -} --- en/03-build-dockerfile-missing-content-json/argv -- -crater image build dockerfile --name img --tag v1 --json --no-interactive --- en/03-build-dockerfile-missing-content-json/exit -- -2 --- en/03-build-dockerfile-missing-content-json/stdout -- --- en/03-build-dockerfile-missing-content-json/stderr -- -{ - "category": "usage_error", - "code": "ERR_MISSING_REQUIRED_FLAG", - "message": "Dockerfile content is required (--dockerfile)" -} --- en/04-build-envd-invalid-source-json/argv -- -crater image build envd --name img --tag v1 --envd x --build-source bad --json --no-interactive --- en/04-build-envd-invalid-source-json/exit -- -2 --- en/04-build-envd-invalid-source-json/stdout -- --- en/04-build-envd-invalid-source-json/stderr -- -{ - "category": "usage_error", - "code": "ERR_INVALID_FLAG_VALUE", - "message": "invalid build-source: bad" -} --- en/05-upload-missing-image-json/argv -- -crater image upload --json --no-interactive --- en/05-upload-missing-image-json/exit -- -2 --- en/05-upload-missing-image-json/stdout -- --- en/05-upload-missing-image-json/stderr -- -{ - "category": "usage_error", - "code": "ERR_MISSING_REQUIRED_FLAG", - "message": "Base image or image link is required (--image)" -} --- en/06-delete-invalid-id-json/argv -- -crater image delete abc --json --no-interactive --- en/06-delete-invalid-id-json/exit -- -2 --- en/06-delete-invalid-id-json/stdout -- --- en/06-delete-invalid-id-json/stderr -- -{ - "category": "usage_error", - "code": "ERR_INVALID_FLAG_VALUE", - "message": "invalid image ID: abc" -} --- en/07-delete-many-invalid-ids-json/argv -- -crater image delete-many --ids 1,x --json --no-interactive --- en/07-delete-many-invalid-ids-json/exit -- -2 --- en/07-delete-many-invalid-ids-json/stdout -- --- en/07-delete-many-invalid-ids-json/stderr -- -{ - "category": "usage_error", - "code": "ERR_INVALID_FLAG_VALUE", - "message": "invalid ID list: 1,x" -} --- en/08-share-add-missing-ids-json/argv -- -crater image share add 1 --json --no-interactive --- en/08-share-add-missing-ids-json/exit -- -2 --- en/08-share-add-missing-ids-json/stdout -- --- en/08-share-add-missing-ids-json/stderr -- -{ - "category": "usage_error", - "code": "ERR_MISSING_REQUIRED_FLAG", - "message": "Comma-separated IDs is required (--ids)" -} --- en/09-share-remove-missing-target-json/argv -- -crater image share remove 1 --json --no-interactive --- en/09-share-remove-missing-target-json/exit -- -2 --- en/09-share-remove-missing-target-json/stdout -- --- en/09-share-remove-missing-target-json/stderr -- -{ - "category": "usage_error", - "code": "ERR_MISSING_REQUIRED_FLAG", - "message": "Share target ID is required (--target-id)" -} --- en/10-harbor-credential-confirm-json/argv -- -crater image harbor credential --json --no-interactive --- en/10-harbor-credential-confirm-json/exit -- -2 --- en/10-harbor-credential-confirm-json/stdout -- --- en/10-harbor-credential-confirm-json/stderr -- -{ - "category": "usage_error", - "code": "ERR_MISSING_REQUIRED_FLAG", - "message": "Force operation without confirmation is required (--yes)" -} --- en/11-cuda-add-missing-json/argv -- -crater image cuda add --json --no-interactive --- en/11-cuda-add-missing-json/exit -- -2 --- en/11-cuda-add-missing-json/stdout -- --- en/11-cuda-add-missing-json/stderr -- -{ - "category": "usage_error", - "code": "ERR_UNKNOWN_COMMAND", - "message": "unknown command \"add\" for \"crater image cuda\"\nRun \"crater image cuda --help\" for usage." -} --- en/12-valid-missing-links-json/argv -- -crater image valid --json --no-interactive --- en/12-valid-missing-links-json/exit -- -2 --- en/12-valid-missing-links-json/stdout -- --- en/12-valid-missing-links-json/stderr -- -{ - "category": "usage_error", - "code": "ERR_MISSING_REQUIRED_FLAG", - "message": "Comma-separated image links is required (--links)" -} --- en/13-build-get-missing-json/argv -- -crater image build get --json --no-interactive --- en/13-build-get-missing-json/exit -- -2 --- en/13-build-get-missing-json/stdout -- --- en/13-build-get-missing-json/stderr -- -{ - "category": "usage_error", - "code": "ERR_MISSING_REQUIRED_FLAG", - "message": "name is required (\u003cname\u003e)" -} --- en/14-build-pod-invalid-json/argv -- -crater image build pod bad --json --no-interactive --- en/14-build-pod-invalid-json/exit -- -2 --- en/14-build-pod-invalid-json/stdout -- --- en/14-build-pod-invalid-json/stderr -- -{ - "category": "usage_error", - "code": "ERR_INVALID_FLAG_VALUE", - "message": "invalid image ID: bad" -} --- en/15-admin-build-remove-missing-json/argv -- -crater admin image build-remove --json --no-interactive --- en/15-admin-build-remove-missing-json/exit -- -2 --- en/15-admin-build-remove-missing-json/stdout -- --- en/15-admin-build-remove-missing-json/stderr -- -{ - "category": "usage_error", - "code": "ERR_MISSING_REQUIRED_FLAG", - "message": "Comma-separated IDs is required (--ids)" -} --- en/16-admin-type-invalid-json/argv -- -crater admin image type 1 --type all --json --no-interactive --- en/16-admin-type-invalid-json/exit -- -2 --- en/16-admin-type-invalid-json/stdout -- --- en/16-admin-type-invalid-json/stderr -- -{ - "category": "usage_error", - "code": "ERR_INVALID_FLAG_VALUE", - "message": "invalid type: all" -} --- en/17-image-ls-404-json/argv -- -crater image ls --json --no-interactive --- en/17-image-ls-404-json/exit -- -4 --- en/17-image-ls-404-json/stdout -- --- en/17-image-ls-404-json/stderr -- -{ - "category": "api_error", - "code": "ERR_NOT_FOUND_404", - "message": "HTTP 404: simulated", - "context": { - "crater_code": 404, - "http_status": 404, - "msg": "simulated" - } -} --- en/18-image-available-404-json/argv -- -crater image ls --available --json --no-interactive --- en/18-image-available-404-json/exit -- -4 --- en/18-image-available-404-json/stdout -- --- en/18-image-available-404-json/stderr -- -{ - "category": "api_error", - "code": "ERR_NOT_FOUND_404", - "message": "HTTP 404: simulated", - "context": { - "crater_code": 404, - "http_status": 404, - "msg": "simulated" - } -} --- en/19-build-pip-apt-404-json/argv -- -crater image build pip-apt --name img --tag v1 --image base:latest --json --no-interactive --- en/19-build-pip-apt-404-json/exit -- -4 --- en/19-build-pip-apt-404-json/stdout -- --- en/19-build-pip-apt-404-json/stderr -- -{ - "category": "api_error", - "code": "ERR_NOT_FOUND_404", - "message": "HTTP 404: simulated", - "context": { - "crater_code": 404, - "http_status": 404, - "msg": "simulated" - } -} --- en/20-share-ls-404-json/argv -- -crater image share ls 1 --json --no-interactive --- en/20-share-ls-404-json/exit -- -4 --- en/20-share-ls-404-json/stdout -- --- en/20-share-ls-404-json/stderr -- -{ - "category": "api_error", - "code": "ERR_NOT_FOUND_404", - "message": "HTTP 404: simulated", - "context": { - "crater_code": 404, - "http_status": 404, - "msg": "simulated" - } -} --- en/21-admin-image-ls-404-json/argv -- -crater admin image ls --json --no-interactive --- en/21-admin-image-ls-404-json/exit -- -4 --- en/21-admin-image-ls-404-json/stdout -- --- en/21-admin-image-ls-404-json/stderr -- -{ - "category": "api_error", - "code": "ERR_NOT_FOUND_404", - "message": "HTTP 404: simulated", - "context": { - "crater_code": 404, - "http_status": 404, - "msg": "simulated" - } -} --- en/22-build-dockerfile-mutually-exclusive-json/argv -- -crater image build dockerfile --name img --tag v1 --dockerfile FROM scratch --file Dockerfile --json --no-interactive --- en/22-build-dockerfile-mutually-exclusive-json/exit -- -2 --- en/22-build-dockerfile-mutually-exclusive-json/stdout -- --- en/22-build-dockerfile-mutually-exclusive-json/stderr -- -{ - "category": "usage_error", - "code": "ERR_INVALID_FLAG_VALUE", - "message": "--dockerfile and --file cannot be used together" -} --- en/23-image-arch-missing-json/argv -- -crater image arch 1 --json --no-interactive --- en/23-image-arch-missing-json/exit -- -2 --- en/23-image-arch-missing-json/stdout -- --- en/23-image-arch-missing-json/stderr -- -{ - "category": "usage_error", - "code": "ERR_MISSING_REQUIRED_FLAG", - "message": "Comma-separated architectures is required (--archs)" -} --- en/24-admin-cuda-add-missing-json/argv -- -crater admin image cuda add --json --no-interactive --- en/24-admin-cuda-add-missing-json/exit -- -2 --- en/24-admin-cuda-add-missing-json/stdout -- --- en/24-admin-cuda-add-missing-json/stderr -- -{ - "category": "usage_error", - "code": "ERR_MISSING_REQUIRED_FLAG", - "message": "Image label is required (--image-label)\nDisplay label is required (--label)\nImage value is required (--value)", - "context": { - "issues": [ - { - "code": "ERR_MISSING_REQUIRED_FLAG", - "field": "image-label", - "message": "Image label is required (--image-label)" - }, - { - "code": "ERR_MISSING_REQUIRED_FLAG", - "field": "label", - "message": "Display label is required (--label)" - }, - { - "code": "ERR_MISSING_REQUIRED_FLAG", - "field": "value", - "message": "Image value is required (--value)" - } - ] - } -} --- en/25-image-help/argv -- -crater image --help --- en/25-image-help/exit -- -0 --- en/25-image-help/stdout -- -Build, upload, delete, share, and update Crater images and image build records. - -Usage: - crater image [flags] - crater image [command] - -Available Commands: - arch Update image architectures - build Manage image builds - cuda View CUDA base images - delete Delete an image - delete-many Delete multiple images - description Update image description - harbor View Harbor information - ls List images - quota View or update Harbor project quota - share Manage image sharing - tags Update image tags - type Update image task type - upload Upload/register an existing image link - valid Validate image links - -Global Flags: - -h, --help Help for crater - --json Output in raw JSON format - --no-interactive Disable interactive prompts - -Use "crater image [command] --help" for more information about a command. --- en/25-image-help/stderr -- --- en/26-image-upload-help/argv -- -crater image upload --help --- en/26-image-upload-help/exit -- -0 --- en/26-image-upload-help/stdout -- -Upload/register an existing image link - -Usage: - crater image upload [flags] - -Flags: - --archs string Comma-separated architectures - --description string Description - --image string Base image or image link - --tags string Comma-separated tags - --type string Image task type (default "custom") - -Global Flags: - -h, --help Help for crater - --json Output in raw JSON format - --no-interactive Disable interactive prompts --- en/26-image-upload-help/stderr -- --- en/27-image-build-pip-apt-help/argv -- -crater image build pip-apt --help --- en/27-image-build-pip-apt-help/exit -- -0 --- en/27-image-build-pip-apt-help/stdout -- -Build an image from base image plus pip/apt packages - -Usage: - crater image build pip-apt [flags] - -Flags: - --archs string Comma-separated architectures - --description string Description - --image string Base image or image link - --name string Image name - --packages string APT packages, separated by spaces - --requirements string Python requirements content - --tag string Image tag - --tags string Comma-separated tags - --template string Template text - -Global Flags: - -h, --help Help for crater - --json Output in raw JSON format - --no-interactive Disable interactive prompts --- en/27-image-build-pip-apt-help/stderr -- diff --git a/cli/testdata/snapshots/image/image.zh-CN.txtar b/cli/testdata/snapshots/image/image.zh-CN.txtar deleted file mode 100644 index 57445c31b..000000000 --- a/cli/testdata/snapshots/image/image.zh-CN.txtar +++ /dev/null @@ -1,405 +0,0 @@ -# Crater CLI snapshot bundle (txtar). Regenerate: make snapshot-update (or UPDATE_SNAPSHOTS=1 go test ./test/snapshots/...) --- zh-CN/01-image-typo-json/argv -- -crater image list --json --no-interactive --- zh-CN/01-image-typo-json/exit -- -2 --- zh-CN/01-image-typo-json/stdout -- --- zh-CN/01-image-typo-json/stderr -- -{ - "category": "usage_error", - "code": "ERR_UNKNOWN_COMMAND", - "message": "unknown command \"list\" for \"crater image\"\n\nDid you mean this?\n\tls\n\nRun \"crater image --help\" for usage." -} --- zh-CN/02-build-pip-apt-missing-json/argv -- -crater image build pip-apt --json --no-interactive --- zh-CN/02-build-pip-apt-missing-json/exit -- -2 --- zh-CN/02-build-pip-apt-missing-json/stdout -- --- zh-CN/02-build-pip-apt-missing-json/stderr -- -{ - "category": "usage_error", - "code": "ERR_MISSING_REQUIRED_FLAG", - "message": "缺少必要参数：镜像名称 (--name)\n缺少必要参数：镜像标签 (--tag)", - "context": { - "issues": [ - { - "code": "ERR_MISSING_REQUIRED_FLAG", - "field": "name", - "message": "缺少必要参数：镜像名称 (--name)" - }, - { - "code": "ERR_MISSING_REQUIRED_FLAG", - "field": "tag", - "message": "缺少必要参数：镜像标签 (--tag)" - } - ] - } -} --- zh-CN/03-build-dockerfile-missing-content-json/argv -- -crater image build dockerfile --name img --tag v1 --json --no-interactive --- zh-CN/03-build-dockerfile-missing-content-json/exit -- -2 --- zh-CN/03-build-dockerfile-missing-content-json/stdout -- --- zh-CN/03-build-dockerfile-missing-content-json/stderr -- -{ - "category": "usage_error", - "code": "ERR_MISSING_REQUIRED_FLAG", - "message": "缺少必要参数：Dockerfile 内容 (--dockerfile)" -} --- zh-CN/04-build-envd-invalid-source-json/argv -- -crater image build envd --name img --tag v1 --envd x --build-source bad --json --no-interactive --- zh-CN/04-build-envd-invalid-source-json/exit -- -2 --- zh-CN/04-build-envd-invalid-source-json/stdout -- --- zh-CN/04-build-envd-invalid-source-json/stderr -- -{ - "category": "usage_error", - "code": "ERR_INVALID_FLAG_VALUE", - "message": "无效的 build-source：bad" -} --- zh-CN/05-upload-missing-image-json/argv -- -crater image upload --json --no-interactive --- zh-CN/05-upload-missing-image-json/exit -- -2 --- zh-CN/05-upload-missing-image-json/stdout -- --- zh-CN/05-upload-missing-image-json/stderr -- -{ - "category": "usage_error", - "code": "ERR_MISSING_REQUIRED_FLAG", - "message": "缺少必要参数：基础镜像或镜像链接 (--image)" -} --- zh-CN/06-delete-invalid-id-json/argv -- -crater image delete abc --json --no-interactive --- zh-CN/06-delete-invalid-id-json/exit -- -2 --- zh-CN/06-delete-invalid-id-json/stdout -- --- zh-CN/06-delete-invalid-id-json/stderr -- -{ - "category": "usage_error", - "code": "ERR_INVALID_FLAG_VALUE", - "message": "无效的镜像 ID：abc" -} --- zh-CN/07-delete-many-invalid-ids-json/argv -- -crater image delete-many --ids 1,x --json --no-interactive --- zh-CN/07-delete-many-invalid-ids-json/exit -- -2 --- zh-CN/07-delete-many-invalid-ids-json/stdout -- --- zh-CN/07-delete-many-invalid-ids-json/stderr -- -{ - "category": "usage_error", - "code": "ERR_INVALID_FLAG_VALUE", - "message": "无效的 ID 列表：1,x" -} --- zh-CN/08-share-add-missing-ids-json/argv -- -crater image share add 1 --json --no-interactive --- zh-CN/08-share-add-missing-ids-json/exit -- -2 --- zh-CN/08-share-add-missing-ids-json/stdout -- --- zh-CN/08-share-add-missing-ids-json/stderr -- -{ - "category": "usage_error", - "code": "ERR_MISSING_REQUIRED_FLAG", - "message": "缺少必要参数：逗号分隔 ID (--ids)" -} --- zh-CN/09-share-remove-missing-target-json/argv -- -crater image share remove 1 --json --no-interactive --- zh-CN/09-share-remove-missing-target-json/exit -- -2 --- zh-CN/09-share-remove-missing-target-json/stdout -- --- zh-CN/09-share-remove-missing-target-json/stderr -- -{ - "category": "usage_error", - "code": "ERR_MISSING_REQUIRED_FLAG", - "message": "缺少必要参数：分享目标 ID (--target-id)" -} --- zh-CN/10-harbor-credential-confirm-json/argv -- -crater image harbor credential --json --no-interactive --- zh-CN/10-harbor-credential-confirm-json/exit -- -2 --- zh-CN/10-harbor-credential-confirm-json/stdout -- --- zh-CN/10-harbor-credential-confirm-json/stderr -- -{ - "category": "usage_error", - "code": "ERR_MISSING_REQUIRED_FLAG", - "message": "缺少必要参数：强制执行操作，无需确认 (--yes)" -} --- zh-CN/11-cuda-add-missing-json/argv -- -crater image cuda add --json --no-interactive --- zh-CN/11-cuda-add-missing-json/exit -- -2 --- zh-CN/11-cuda-add-missing-json/stdout -- --- zh-CN/11-cuda-add-missing-json/stderr -- -{ - "category": "usage_error", - "code": "ERR_UNKNOWN_COMMAND", - "message": "unknown command \"add\" for \"crater image cuda\"\nRun \"crater image cuda --help\" for usage." -} --- zh-CN/12-valid-missing-links-json/argv -- -crater image valid --json --no-interactive --- zh-CN/12-valid-missing-links-json/exit -- -2 --- zh-CN/12-valid-missing-links-json/stdout -- --- zh-CN/12-valid-missing-links-json/stderr -- -{ - "category": "usage_error", - "code": "ERR_MISSING_REQUIRED_FLAG", - "message": "缺少必要参数：逗号分隔镜像链接 (--links)" -} --- zh-CN/13-build-get-missing-json/argv -- -crater image build get --json --no-interactive --- zh-CN/13-build-get-missing-json/exit -- -2 --- zh-CN/13-build-get-missing-json/stdout -- --- zh-CN/13-build-get-missing-json/stderr -- -{ - "category": "usage_error", - "code": "ERR_MISSING_REQUIRED_FLAG", - "message": "缺少必要参数：name (\u003cname\u003e)" -} --- zh-CN/14-build-pod-invalid-json/argv -- -crater image build pod bad --json --no-interactive --- zh-CN/14-build-pod-invalid-json/exit -- -2 --- zh-CN/14-build-pod-invalid-json/stdout -- --- zh-CN/14-build-pod-invalid-json/stderr -- -{ - "category": "usage_error", - "code": "ERR_INVALID_FLAG_VALUE", - "message": "无效的镜像 ID：bad" -} --- zh-CN/15-admin-build-remove-missing-json/argv -- -crater admin image build-remove --json --no-interactive --- zh-CN/15-admin-build-remove-missing-json/exit -- -2 --- zh-CN/15-admin-build-remove-missing-json/stdout -- --- zh-CN/15-admin-build-remove-missing-json/stderr -- -{ - "category": "usage_error", - "code": "ERR_MISSING_REQUIRED_FLAG", - "message": "缺少必要参数：逗号分隔 ID (--ids)" -} --- zh-CN/16-admin-type-invalid-json/argv -- -crater admin image type 1 --type all --json --no-interactive --- zh-CN/16-admin-type-invalid-json/exit -- -2 --- zh-CN/16-admin-type-invalid-json/stdout -- --- zh-CN/16-admin-type-invalid-json/stderr -- -{ - "category": "usage_error", - "code": "ERR_INVALID_FLAG_VALUE", - "message": "无效的 type：all" -} --- zh-CN/17-image-ls-404-json/argv -- -crater image ls --json --no-interactive --- zh-CN/17-image-ls-404-json/exit -- -4 --- zh-CN/17-image-ls-404-json/stdout -- --- zh-CN/17-image-ls-404-json/stderr -- -{ - "category": "api_error", - "code": "ERR_NOT_FOUND_404", - "message": "请求失败（HTTP 404）：simulated", - "context": { - "crater_code": 404, - "http_status": 404, - "msg": "simulated" - } -} --- zh-CN/18-image-available-404-json/argv -- -crater image ls --available --json --no-interactive --- zh-CN/18-image-available-404-json/exit -- -4 --- zh-CN/18-image-available-404-json/stdout -- --- zh-CN/18-image-available-404-json/stderr -- -{ - "category": "api_error", - "code": "ERR_NOT_FOUND_404", - "message": "请求失败（HTTP 404）：simulated", - "context": { - "crater_code": 404, - "http_status": 404, - "msg": "simulated" - } -} --- zh-CN/19-build-pip-apt-404-json/argv -- -crater image build pip-apt --name img --tag v1 --image base:latest --json --no-interactive --- zh-CN/19-build-pip-apt-404-json/exit -- -4 --- zh-CN/19-build-pip-apt-404-json/stdout -- --- zh-CN/19-build-pip-apt-404-json/stderr -- -{ - "category": "api_error", - "code": "ERR_NOT_FOUND_404", - "message": "请求失败（HTTP 404）：simulated", - "context": { - "crater_code": 404, - "http_status": 404, - "msg": "simulated" - } -} --- zh-CN/20-share-ls-404-json/argv -- -crater image share ls 1 --json --no-interactive --- zh-CN/20-share-ls-404-json/exit -- -4 --- zh-CN/20-share-ls-404-json/stdout -- --- zh-CN/20-share-ls-404-json/stderr -- -{ - "category": "api_error", - "code": "ERR_NOT_FOUND_404", - "message": "请求失败（HTTP 404）：simulated", - "context": { - "crater_code": 404, - "http_status": 404, - "msg": "simulated" - } -} --- zh-CN/21-admin-image-ls-404-json/argv -- -crater admin image ls --json --no-interactive --- zh-CN/21-admin-image-ls-404-json/exit -- -4 --- zh-CN/21-admin-image-ls-404-json/stdout -- --- zh-CN/21-admin-image-ls-404-json/stderr -- -{ - "category": "api_error", - "code": "ERR_NOT_FOUND_404", - "message": "请求失败（HTTP 404）：simulated", - "context": { - "crater_code": 404, - "http_status": 404, - "msg": "simulated" - } -} --- zh-CN/22-build-dockerfile-mutually-exclusive-json/argv -- -crater image build dockerfile --name img --tag v1 --dockerfile FROM scratch --file Dockerfile --json --no-interactive --- zh-CN/22-build-dockerfile-mutually-exclusive-json/exit -- -2 --- zh-CN/22-build-dockerfile-mutually-exclusive-json/stdout -- --- zh-CN/22-build-dockerfile-mutually-exclusive-json/stderr -- -{ - "category": "usage_error", - "code": "ERR_INVALID_FLAG_VALUE", - "message": "--dockerfile 和 --file 不能同时使用" -} --- zh-CN/23-image-arch-missing-json/argv -- -crater image arch 1 --json --no-interactive --- zh-CN/23-image-arch-missing-json/exit -- -2 --- zh-CN/23-image-arch-missing-json/stdout -- --- zh-CN/23-image-arch-missing-json/stderr -- -{ - "category": "usage_error", - "code": "ERR_MISSING_REQUIRED_FLAG", - "message": "缺少必要参数：逗号分隔架构 (--archs)" -} --- zh-CN/24-admin-cuda-add-missing-json/argv -- -crater admin image cuda add --json --no-interactive --- zh-CN/24-admin-cuda-add-missing-json/exit -- -2 --- zh-CN/24-admin-cuda-add-missing-json/stdout -- --- zh-CN/24-admin-cuda-add-missing-json/stderr -- -{ - "category": "usage_error", - "code": "ERR_MISSING_REQUIRED_FLAG", - "message": "缺少必要参数：镜像标签 (--image-label)\n缺少必要参数：展示标签 (--label)\n缺少必要参数：镜像值 (--value)", - "context": { - "issues": [ - { - "code": "ERR_MISSING_REQUIRED_FLAG", - "field": "image-label", - "message": "缺少必要参数：镜像标签 (--image-label)" - }, - { - "code": "ERR_MISSING_REQUIRED_FLAG", - "field": "label", - "message": "缺少必要参数：展示标签 (--label)" - }, - { - "code": "ERR_MISSING_REQUIRED_FLAG", - "field": "value", - "message": "缺少必要参数：镜像值 (--value)" - } - ] - } -} --- zh-CN/25-image-help/argv -- -crater image --help --- zh-CN/25-image-help/exit -- -0 --- zh-CN/25-image-help/stdout -- -构建、上传、删除、分享和更新 Crater 镜像与镜像构建记录。 - -Usage: - crater image [flags] - crater image [command] - -Available Commands: - arch 更新镜像架构 - build 管理镜像构建 - cuda 查看 CUDA 基础镜像 - delete 删除镜像 - delete-many 批量删除镜像 - description 更新镜像描述 - harbor 查看 Harbor 信息 - ls 列出镜像 - quota 查看或更新 Harbor 项目配额 - share 管理镜像分享 - tags 更新镜像标签 - type 更新镜像任务类型 - upload 上传/登记已有镜像链接 - valid 校验镜像链接 - -Global Flags: - -h, --help 显示帮助信息 - --json 以原始 JSON 格式输出 - --no-interactive 禁用交互式提示 - -Use "crater image [command] --help" for more information about a command. --- zh-CN/25-image-help/stderr -- --- zh-CN/26-image-upload-help/argv -- -crater image upload --help --- zh-CN/26-image-upload-help/exit -- -0 --- zh-CN/26-image-upload-help/stdout -- -上传/登记已有镜像链接 - -Usage: - crater image upload [flags] - -Flags: - --archs string 逗号分隔架构 - --description string 描述 - --image string 基础镜像或镜像链接 - --tags string 逗号分隔标签 - --type string 镜像任务类型 (default "custom") - -Global Flags: - -h, --help 显示帮助信息 - --json 以原始 JSON 格式输出 - --no-interactive 禁用交互式提示 --- zh-CN/26-image-upload-help/stderr -- --- zh-CN/27-image-build-pip-apt-help/argv -- -crater image build pip-apt --help --- zh-CN/27-image-build-pip-apt-help/exit -- -0 --- zh-CN/27-image-build-pip-apt-help/stdout -- -基于基础镜像和 pip/apt 包构建镜像 - -Usage: - crater image build pip-apt [flags] - -Flags: - --archs string 逗号分隔架构 - --description string 描述 - --image string 基础镜像或镜像链接 - --name string 镜像名称 - --packages string APT 包，使用空格分隔 - --requirements string Python requirements 内容 - --tag string 镜像标签 - --tags string 逗号分隔标签 - --template string 模板文本 - -Global Flags: - -h, --help 显示帮助信息 - --json 以原始 JSON 格式输出 - --no-interactive 禁用交互式提示 --- zh-CN/27-image-build-pip-apt-help/stderr -- diff --git a/crater-agent/.dockerignore b/crater-agent/.dockerignore new file mode 100644 index 000000000..8f839ed36 --- /dev/null +++ b/crater-agent/.dockerignore @@ -0,0 +1,10 @@ +.venv +__pycache__ +*.pyc +*.pyo +.pytest_cache +.ruff_cache +logs +.env +config/kubeconfig +config/*backup* diff --git a/crater-agent/.gitignore b/crater-agent/.gitignore new file mode 100644 index 000000000..20e1fad11 --- /dev/null +++ b/crater-agent/.gitignore @@ -0,0 +1,34 @@ +__pycache__/ +*.py[cod] +*.egg-info/ +dist/ +.venv/ +.env +*.log +.ruff_cache/ +.pytest_cache/ + +# Offline evaluation / local experiment artifacts +/dataset/ +/crater_bench/ +/results/ +/logs/crater-agent/eval/ +/crater_agent/dataset/ +/crater_agent/eval/ +/crater_agent/res/ + +# Local-only runtime configuration +/config/kubeconfig +/config/*offline*.json +/config/*offline*.yaml +/config/*bench*.json +/config/*bench*.yaml +/config/*tune*.json +/config/*probe*.json +/config/*plan-execute*.json +/config/mops*.json + +# Large local dumps +*.jsonl +*.tar +*.zip diff --git a/crater-agent/Dockerfile b/crater-agent/Dockerfile new file mode 100644 index 000000000..11f1cb553 --- /dev/null +++ b/crater-agent/Dockerfile @@ -0,0 +1,37 @@ +# Copyright 2025 RAIDS Lab +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +FROM python:3.11-slim + +LABEL org.opencontainers.image.source="https://github.com/raids-lab/crater-agent" +LABEL org.opencontainers.image.description="Crater AI operations agent service" +LABEL org.opencontainers.image.licenses="Apache-2.0" + +ENV PYTHONDONTWRITEBYTECODE=1 +ENV PYTHONUNBUFFERED=1 +ENV CRATER_AGENT_HOST=0.0.0.0 +ENV CRATER_AGENT_PORT=8000 + +WORKDIR /app + +COPY pyproject.toml README.md ./ +COPY config ./config +COPY crater_agent ./crater_agent + +RUN python -m pip install --no-cache-dir --upgrade pip \ + && python -m pip install --no-cache-dir . + +EXPOSE 8000 + +CMD ["sh", "-c", "python -m uvicorn crater_agent.app:app --host ${CRATER_AGENT_HOST:-0.0.0.0} --port ${CRATER_AGENT_PORT:-8000}"] diff --git a/crater-agent/README.md b/crater-agent/README.md new file mode 100644 index 000000000..50bb5a562 --- /dev/null +++ b/crater-agent/README.md @@ -0,0 +1,5 @@ +# Crater Agent + +Crater Agent is the Python service used by the Crater AI assistant. It receives chat requests from the Go backend, calls platform tools through internal backend APIs, and streams agent responses back to the browser. + +The service is packaged as a container image and deployed by the Crater Helm chart when `agent.enabled` is true. diff --git a/crater-agent/config/llm-clients.json b/crater-agent/config/llm-clients.json new file mode 100644 index 000000000..1f244fcd2 --- /dev/null +++ b/crater-agent/config/llm-clients.json @@ -0,0 +1,13 @@ +{ + "default": { + "provider": "openai_compatible", + "base_url": "", + "api_key_env": "", + "model": "", + "temperature": 0.1, + "max_tokens": 8192, + "timeout": 120, + "streaming": true, + "stream_usage": true + } +} diff --git a/crater-agent/crater_agent/__init__.py b/crater-agent/crater_agent/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/crater-agent/crater_agent/agent/__init__.py b/crater-agent/crater_agent/agent/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/crater-agent/crater_agent/agent/compaction.py b/crater-agent/crater_agent/agent/compaction.py new file mode 100644 index 000000000..c0e93a69a --- /dev/null +++ b/crater-agent/crater_agent/agent/compaction.py @@ -0,0 +1,173 @@ +"""LLM-based conversation history compaction. + +When the context window is nearing its limit, this module summarises older +conversation messages using the LLM, preserving key investigation context +(tool calls, findings, conclusions) while drastically reducing token count. + +Design follows HolmesGPT's compaction approach: +- Tool schemas (bind_tools) are NEVER compressed — they're a separate API param. +- SystemMessage is always preserved uncompacted. +- The most recent UserMessage and the last N messages are preserved. +- Everything else is summarised into a single AIMessage. +""" + +from __future__ import annotations + +import asyncio +import logging +from typing import Any + +from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage + +logger = logging.getLogger(__name__) + +COMPACTION_PROMPT = """\ +请将以下对话历史压缩为简洁的中文摘要。 + +## 当前用户正在问的问题： +{current_query} + +## 历史对话中的用户消息： +{user_messages} + +## 历史对话中的助手回复和工具调用结果： +{compressible_messages} + +## 压缩规则： +1. **与当前问题相关的历史**：保留关键结论、工具名+参数+结果、技术细节（资源名/命名空间/错误码） +2. **与当前问题无关的历史**：一句话概括即可（如"之前用户询问过 job-X 的状态，已解决"），不需要保留细节 +3. 用户的原始意图用简短引用保留，不要改写用户的核心诉求 +4. 助手的冗长分析压缩为 1-2 句结论 +5. 如果有待完成的操作或未解决的问题，必须保留 + +注意：历史消息不一定都和当前问题相关，请根据相关性灵活压缩，不需要平均对待每条历史。 + +请用中文输出紧凑摘要。""" + + +def _format_messages_for_summary(messages: list[Any]) -> tuple[str, str]: + """Split messages into user messages (preserve) and compressible messages. + + Returns: + (user_messages_text, compressible_messages_text) + """ + user_parts: list[str] = [] + compress_parts: list[str] = [] + for msg in messages: + content = str(getattr(msg, "content", "") or "") + if not content: + continue + if isinstance(msg, HumanMessage): + user_parts.append(f"- {content}") + elif isinstance(msg, AIMessage): + tool_calls = getattr(msg, "tool_calls", None) or [] + if tool_calls: + tc_names = ", ".join(tc.get("name", "?") for tc in tool_calls) + compress_parts.append(f"[助手→调用工具: {tc_names}] {content}") + else: + compress_parts.append(f"[助手] {content}") + elif isinstance(msg, ToolMessage): + tool_call_id = getattr(msg, "tool_call_id", "") + compress_parts.append(f"[工具结果 {tool_call_id}] {content}") + elif isinstance(msg, SystemMessage): + continue + else: + compress_parts.append(f"[{type(msg).__name__}] {content}") + return "\n".join(user_parts) or "（无）", "\n".join(compress_parts) or "（无）" + + +async def compact_messages_with_llm( + messages: list[Any], + llm: Any, + preserve_tail: int = 4, + timeout: float = 15.0, + current_query: str = "", +) -> list[Any] | None: + """Summarise older messages using the LLM to reduce context window usage. + + Args: + messages: Full message list (including SystemMessage at index 0). + llm: A ChatOpenAI (or compatible) LLM instance. + preserve_tail: Number of most-recent messages to keep uncompacted. + timeout: Maximum seconds to wait for the LLM summary call. + current_query: The current user question — used as anchor for + relevance-based compression. History unrelated to this + query is compressed more aggressively. + + Returns: + Compacted message list on success, or ``None`` on failure (caller + should fall back to the existing hard-truncation strategy). + """ + if len(messages) <= preserve_tail + 2: + return None # Not enough messages to compact + + # Split: system message + body + system_msg = messages[0] if isinstance(messages[0], SystemMessage) else None + body = messages[1:] if system_msg else list(messages) + + if len(body) <= preserve_tail: + return None + + # Partition into compactable and preserved sections + # We need to keep at least `preserve_tail` messages + start_idx = max(0, len(body) - preserve_tail) + + # CRITICAL: Never start the tail with a ToolMessage. + # If the cutoff point lands on a ToolMessage, we must include the preceding AIMessage + # that made the tool_calls, otherwise the LLM API will reject the sequence. + from langchain_core.messages import ToolMessage + while start_idx > 0 and isinstance(body[start_idx], ToolMessage): + start_idx -= 1 + + to_compact = body[:start_idx] + tail = body[start_idx:] + + # Ensure the last user message is in the preserved section + last_human = next( + (m for m in reversed(body) if isinstance(m, HumanMessage)), None + ) + tail_has_human = any(isinstance(m, HumanMessage) for m in tail) + + conversation_text = _format_messages_for_summary(to_compact) + user_text, compress_text = conversation_text + if not user_text.strip() and not compress_text.strip(): + return None + + # Derive current query from preserved tail if not explicitly provided + if not current_query: + for m in reversed(tail): + if isinstance(m, HumanMessage): + current_query = str(m.content or "")[:300] + break + + prompt = COMPACTION_PROMPT.format( + current_query=current_query or "（未提供）", + user_messages=user_text, + compressible_messages=compress_text, + ) + + try: + response = await asyncio.wait_for( + llm.ainvoke([HumanMessage(content=prompt)]), + timeout=timeout, + ) + summary = str(getattr(response, "content", "") or "").strip() + if not summary: + logger.warning("LLM compaction returned empty summary, falling back") + return None + except asyncio.TimeoutError: + logger.warning("LLM compaction timed out after %.1fs, falling back", timeout) + return None + except Exception: + logger.warning("LLM compaction failed, falling back", exc_info=True) + return None + + # Assemble compacted messages + compacted: list[Any] = [] + if system_msg: + compacted.append(system_msg) + compacted.append(AIMessage(content=f"[对话摘要] {summary}")) + if last_human and not tail_has_human: + compacted.append(last_human) + compacted.extend(tail) + return compacted diff --git a/crater-agent/crater_agent/agent/graph.py b/crater-agent/crater_agent/agent/graph.py new file mode 100644 index 000000000..b6fe1f97f --- /dev/null +++ b/crater-agent/crater_agent/agent/graph.py @@ -0,0 +1,739 @@ +"""LangGraph ReAct agent graph for Crater. + +The agent uses a simple ReAct loop: + agent (LLM think) → tools (execute) → agent (observe & think again) → ... + until LLM decides to respond without tool calls → END + +Key design: LLM autonomously decides which tools to call and when to stop. +No fixed workflow or intent classification — pure ReAct. +""" + +from __future__ import annotations + +import json +import logging +import re +import time +from typing import Any + +from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage +from langchain_openai import ChatOpenAI +from langgraph.graph import END, StateGraph +from openai import BadRequestError + +from crater_agent.agent.prompts import build_system_prompt +from crater_agent.agent.state import CraterAgentState +from crater_agent.config import settings +from crater_agent.llm.client import ModelClientFactory +from crater_agent.tools.definitions import ALL_TOOLS +from crater_agent.tools.executor import GoBackendToolExecutor, ToolExecutorProtocol +from crater_agent.tools.tool_selector import ( + _resolve_actor_role, + sanitize_capabilities_for_context, + select_tools_for_context, +) + +logger = logging.getLogger(__name__) + + +def _truncate_text(value: str, max_chars: int = 2400) -> str: + if len(value) <= max_chars: + return value + head = value[: max_chars // 2] + tail = value[-max_chars // 2 :] + return f"{head}\n\n...(内容过长，已截断)...\n\n{tail}" + + +_DEFAULT_TOOL_TOKEN_BUDGET = 3000 + +_TOOL_TOKEN_BUDGETS: dict[str, int] = { + "get_job_logs": 4000, + "diagnose_job": 4000, + "get_diagnostic_context": 4000, + "get_job_detail": 3000, + "prometheus_query": 2000, + "query_job_metrics": 2000, +} + +_TOOL_EXTRACT_PROMPT = """\ +从以下工具输出中提取与用户问题最相关的关键信息。 + +工具: {tool_name} +用户问题: {user_question} + +工具完整输出: +{tool_output} + +要求： +- 保留所有错误信息、异常堆栈、关键状态码 +- 保留与用户问题直接相关的数据 +- 保留资源名、命名空间、时间戳等关键标识 +- 删除重复的正常日志行、冗余的健康检查输出 +- 用紧凑格式输出，不要添加额外解释""" + +_JOB_MUTATION_TOOL_BY_INTENT: tuple[tuple[str, tuple[str, ...]], ...] = ( + ("delete_job", ("删除", "删掉", "移除", "delete")), + ("stop_job", ("停止", "终止", "停掉", "stop", "kill")), + ("resubmit_job", ("重新提交", "重提", "再提交", "resubmit")), +) + +_SYSTEM_JOB_NAME_PATTERN = re.compile( + r"\b(?:jpt|webide|custom|pytorch|tensorflow|tf|pt|sg|ds|mpi)-[A-Za-z0-9][A-Za-z0-9_-]*\b", + re.IGNORECASE, +) + + +def _latest_human_text(messages: list[Any]) -> str: + for message in reversed(messages): + if isinstance(message, HumanMessage): + return str(message.content or "") + return "" + + +def _detect_requested_job_mutation_tool(user_text: str) -> str: + normalized = user_text.strip().lower() + for tool_name, keywords in _JOB_MUTATION_TOOL_BY_INTENT: + if any(keyword in normalized for keyword in keywords): + return tool_name + return "" + + +def _extract_system_job_name(text: str) -> str: + matches = _SYSTEM_JOB_NAME_PATTERN.findall(text or "") + if not matches: + return "" + return max(matches, key=len) + + +def _looks_like_text_confirmation(content: str) -> bool: + normalized = str(content or "").strip().lower() + if not normalized: + return False + confirmation_terms = ("确认", "是否", "要删除", "要停止", "要重新提交", "请确认") + return any(term in normalized for term in confirmation_terms) + + +def _coerce_text_confirmation_to_tool_call( + messages: list[Any], + response: AIMessage, +) -> AIMessage: + """Convert textual confirmation prompts into real confirmation-card tools. + + The backend confirmation card is the only valid approval UI. If the model + has already resolved a job mutation target but asks the user to type + "确认", force the matching write tool in the same turn. + """ + if response.tool_calls or not _looks_like_text_confirmation(str(response.content or "")): + return response + + tool_name = _detect_requested_job_mutation_tool(_latest_human_text(messages)) + if not tool_name: + return response + + job_name = _extract_system_job_name(str(response.content or "")) + if not job_name: + for message in reversed(messages): + if isinstance(message, ToolMessage): + job_name = _extract_system_job_name(str(message.content or "")) + if job_name: + break + if not job_name: + return response + + logger.info( + "coercing textual confirmation into tool call: tool=%s job_name=%s", + tool_name, + job_name, + ) + return AIMessage( + content="", + tool_calls=[ + { + "name": tool_name, + "args": {"job_name": job_name}, + "id": f"forced-{tool_name}-{int(time.time() * 1000)}", + } + ], + ) + + +def _truncate_to_token_budget(text: str, budget_tokens: int) -> str: + """Truncate text to fit within a token budget, keeping head + tail.""" + from crater_agent.llm.tokenizer import count_tokens + + if count_tokens(text) <= budget_tokens: + return text + max_chars = budget_tokens * 3 # conservative: ~3 chars/token + return _truncate_text(text, max_chars=max_chars) + + +async def _extract_with_llm( + tool_name: str, + raw_text: str, + budget_tokens: int, + llm: Any, + user_question: str, +) -> str | None: + """Use LLM to intelligently extract key info from oversized tool results. + + Returns extracted text on success, None on failure (caller falls back to + hard truncation). + """ + import asyncio + + prompt = _TOOL_EXTRACT_PROMPT.format( + tool_name=tool_name, + user_question=user_question or "（未知）", + tool_output=raw_text, + ) + try: + response = await asyncio.wait_for( + llm.ainvoke([HumanMessage(content=prompt)]), + timeout=10.0, + ) + extracted = str(getattr(response, "content", "") or "").strip() + if not extracted: + return None + logger.info( + "LLM extract for %s: %d chars -> %d chars", + tool_name, len(raw_text), len(extracted), + ) + return extracted + except asyncio.TimeoutError: + logger.warning("LLM extract for %s timed out", tool_name) + return None + except Exception: + logger.warning("LLM extract for %s failed", tool_name, exc_info=True) + return None + + +def _build_tool_observation_sync(tool_name: str, result: dict[str, Any]) -> str: + """Build tool observation string (synchronous, hard truncation only).""" + if result.get("status") == "confirmation_required": + return json.dumps(result, ensure_ascii=False) + + budget = _TOOL_TOKEN_BUDGETS.get(tool_name, _DEFAULT_TOOL_TOKEN_BUDGET) + + if result.get("status") == "error": + error_payload = { + "status": "error", + "tool_name": tool_name, + "message": result.get("message", ""), + "error_type": result.get("error_type", "unknown"), + "retryable": bool(result.get("retryable", False)), + } + if "status_code" in result: + error_payload["status_code"] = result["status_code"] + if isinstance(result.get("result"), dict): + error_payload["result"] = result["result"] + return _truncate_to_token_budget( + json.dumps(error_payload, ensure_ascii=False), budget + ) + + result_content = result.get("result", result.get("message", "")) + if isinstance(result_content, dict): + return _truncate_to_token_budget( + json.dumps(result_content, ensure_ascii=False), budget + ) + + return _truncate_to_token_budget(str(result_content), budget) + + +async def _build_tool_observation( + tool_name: str, + result: dict[str, Any], + llm: Any = None, + user_question: str = "", +) -> str: + """Build tool observation, using LLM extract for oversized results. + + Flow: + 1. Serialize the raw result + 2. If within token budget → return as-is + 3. If over budget AND llm provided → LLM extract (semantic compression) + 4. If LLM extract fails or no llm → fallback to head+tail hard truncation + """ + if result.get("status") == "confirmation_required": + return json.dumps(result, ensure_ascii=False) + + budget = _TOOL_TOKEN_BUDGETS.get(tool_name, _DEFAULT_TOOL_TOKEN_BUDGET) + from crater_agent.llm.tokenizer import count_tokens + + # Serialize the raw result + if result.get("status") == "error": + error_payload = { + "status": "error", + "tool_name": tool_name, + "message": result.get("message", ""), + "error_type": result.get("error_type", "unknown"), + "retryable": bool(result.get("retryable", False)), + } + if "status_code" in result: + error_payload["status_code"] = result["status_code"] + if isinstance(result.get("result"), dict): + error_payload["result"] = result["result"] + raw_text = json.dumps(error_payload, ensure_ascii=False) + else: + result_content = result.get("result", result.get("message", "")) + if isinstance(result_content, dict): + raw_text = json.dumps(result_content, ensure_ascii=False) + else: + raw_text = str(result_content) + + # Fast path: within budget, return as-is + if count_tokens(raw_text) <= budget: + return raw_text + + # Slow path: over budget → try LLM extract, then fallback to hard truncation + if llm is not None: + extracted = await _extract_with_llm( + tool_name, raw_text, budget, llm, user_question, + ) + if extracted is not None: + # Ensure extracted result fits the budget (LLM may still overshoot) + return _truncate_to_token_budget(extracted, budget) + + return _truncate_to_token_budget(raw_text, budget) + + +def _is_context_limit_error(exc: Exception) -> bool: + message = str(exc or "") + return ( + "exceed_context_size_error" in message + or "available context size" in message + or "maximum context length" in message + ) + + +def _compact_message(message: Any) -> Any: + if isinstance(message, SystemMessage): + return SystemMessage(content=_truncate_text(str(message.content or ""), max_chars=1600)) + if isinstance(message, HumanMessage): + return HumanMessage(content=_truncate_text(str(message.content or ""), max_chars=600)) + if isinstance(message, ToolMessage): + return ToolMessage( + content=_truncate_text(str(message.content or ""), max_chars=800), + tool_call_id=getattr(message, "tool_call_id", "unknown"), + ) + if isinstance(message, AIMessage): + return AIMessage( + content=_truncate_text(str(message.content or ""), max_chars=600), + tool_calls=list(getattr(message, "tool_calls", []) or []), + additional_kwargs=dict(getattr(message, "additional_kwargs", {}) or {}), + response_metadata=dict(getattr(message, "response_metadata", {}) or {}), + ) + return message + + +def _compact_messages_for_retry(messages: list[Any]) -> list[Any]: + if not messages: + return messages + + system_message = messages[0] if isinstance(messages[0], SystemMessage) else None + body = list(messages[1:] if system_message else messages) + tail = body[-6:] + last_human = next((msg for msg in reversed(body) if isinstance(msg, HumanMessage)), None) + + compacted: list[Any] = [] + if system_message: + compacted.append(_compact_message(system_message)) + if last_human and last_human not in tail: + compacted.append(_compact_message(last_human)) + compacted.extend(_compact_message(msg) for msg in tail) + return compacted + + +def _estimate_message_tokens(messages: list[Any]) -> int: + """Count tokens across messages using tiktoken (with heuristic fallback).""" + from crater_agent.llm.tokenizer import count_message_tokens + + return count_message_tokens(messages) + + +def _extract_current_query(messages: list[Any]) -> str: + """Find the most recent user question from the message list.""" + for msg in reversed(messages): + if isinstance(msg, HumanMessage): + return str(msg.content or "")[:300] + return "" + + +async def _proactive_compact( + messages: list[Any], max_context: int, llm: Any = None, +) -> list[Any]: + """Proactively compact messages before hitting the API context limit. + + Reserves budget for tool schemas (~8000 tokens) and LLM response (~4000 tokens). + First attempts LLM-based summarisation; falls back to hard truncation on failure. + """ + tool_schema_budget = 8000 + response_budget = 4000 + available = max_context - tool_schema_budget - response_budget + if available <= 0: + return messages + estimated = _estimate_message_tokens(messages) + if estimated <= available: + return messages + logger.info( + "Proactive compaction: estimated %d tokens > available %d, compacting", + estimated, available, + ) + # Try LLM summarisation first + if llm is not None: + from crater_agent.agent.compaction import compact_messages_with_llm + + compacted = await compact_messages_with_llm( + messages, llm, current_query=_extract_current_query(messages), + ) + if compacted is not None: + logger.info( + "LLM compaction succeeded: %d -> %d messages", + len(messages), len(compacted), + ) + return compacted + # Fallback to hard truncation + return _compact_messages_for_retry(messages) + + +def create_llm() -> ChatOpenAI: + """Create the default single-agent LLM instance.""" + return ModelClientFactory().create("default") + + +def create_agent_graph( + tool_executor: ToolExecutorProtocol | None = None, + llm: ChatOpenAI | None = None, +) -> StateGraph: + """Build the LangGraph StateGraph for the Crater ReAct agent. + + Args: + tool_executor: Tool executor instance. Defaults to GoBackendToolExecutor. + """ + if tool_executor is None: + tool_executor = GoBackendToolExecutor() + + llm = llm or create_llm() + + def get_enabled_tools(context: dict[str, Any]) -> list[Any]: + capabilities = sanitize_capabilities_for_context(context, context.get("capabilities")) + enabled_tool_names = capabilities.get("enabled_tools") or [] + if enabled_tool_names: + enabled_set = set(enabled_tool_names) + base_tools = [tool for tool in ALL_TOOLS if tool.name in enabled_set] + else: + base_tools = ALL_TOOLS + return select_tools_for_context(context, base_tools) + + # ----- Node: agent (LLM reasoning) ----- + async def agent_node(state: CraterAgentState) -> dict: + """LLM thinks about the current state and decides next action.""" + messages = state["messages"] + context = dict(state.get("context", {}) or {}) + context["capabilities"] = sanitize_capabilities_for_context(context, context.get("capabilities")) + llm_with_tools = llm.bind_tools(get_enabled_tools(context)) + + # Build system prompt on first call (check if system message exists) + if not messages or not isinstance(messages[0], SystemMessage): + actor = context.get("actor", {}) + is_first_time = actor.get("is_first_time", False) + current_query = str(messages[-1].content) if messages else "" + system_prompt = build_system_prompt( + context=context, + is_first_time=is_first_time, + user_message=current_query, + ) + messages = [SystemMessage(content=system_prompt)] + list(messages) + + try: + messages = await _proactive_compact( + messages, max_context=settings.max_context_tokens, llm=llm, + ) + llm_start = time.time() + response = await llm_with_tools.ainvoke(messages) + llm_ms = int((time.time() - llm_start) * 1000) + logger.info( + "agent_node: LLM responded in %dms, tool_calls=%d, content_len=%d", + llm_ms, len(response.tool_calls) if response.tool_calls else 0, + len(response.content) if response.content else 0, + ) + except BadRequestError as exc: + if not _is_context_limit_error(exc): + raise + # Try LLM compaction on context limit error before hard fallback + from crater_agent.agent.compaction import compact_messages_with_llm + + compact_messages = await compact_messages_with_llm( + messages, llm, current_query=_extract_current_query(messages), + ) + if compact_messages is None: + compact_messages = _compact_messages_for_retry(messages) + logger.warning( + "Single-agent context limit hit, retrying with compacted messages: %d -> %d", + len(messages), + len(compact_messages), + ) + response = await llm_with_tools.ainvoke(compact_messages) + + # Handle qwen thinking mode: content may be empty, actual reply in reasoning_content + if not response.content and not response.tool_calls: + reasoning = getattr(response, "reasoning_content", "") or ( + response.additional_kwargs or {} + ).get("reasoning_content", "") + if reasoning: + response = AIMessage( + content=reasoning, + additional_kwargs=response.additional_kwargs, + response_metadata=response.response_metadata, + ) + + response = _coerce_text_confirmation_to_tool_call(messages, response) + + # Record trace + trace_entry = { + "node": "agent", + "timestamp": time.time(), + "has_tool_calls": bool(response.tool_calls), + "tool_calls_count": len(response.tool_calls) if response.tool_calls else 0, + "response_length": len(response.content) if response.content else 0, + } + + return { + "messages": [response], + "trace": [trace_entry], + } + + # ----- Node: tools (execute tool calls) ----- + async def tools_node(state: CraterAgentState) -> dict: + """Execute the tool calls from the last AI message.""" + messages = state["messages"] + last_message = messages[-1] + context = dict(state.get("context", {}) or {}) + context["capabilities"] = sanitize_capabilities_for_context(context, context.get("capabilities")) + # Extract user question for LLM tool result extraction context + _user_question = "" + for msg in reversed(messages): + if isinstance(msg, HumanMessage): + _user_question = str(msg.content or "")[:200] + break + session_id = context.get("session_id", "unknown") + actor = context.get("actor", {}) + user_id = actor.get("user_id", 0) + actor_role = _resolve_actor_role(context) + turn_id = context.get("turn_id") + + tool_messages = [] + trace_entries = [] + new_tool_call_count = state.get("tool_call_count", 0) + pending_confirmations: list[dict[str, Any]] = [] + attempted_tool_calls = dict(state.get("attempted_tool_calls") or {}) + + for tc in last_message.tool_calls: + tool_name = tc["name"] + tool_args = tc["args"] + new_tool_call_count += 1 + tool_signature = json.dumps( + { + "tool_name": tool_name, + "tool_args": tool_args, + }, + ensure_ascii=False, + sort_keys=True, + ) + if attempted_tool_calls.get(tool_signature, 0) >= 1: + tool_messages.append( + ToolMessage( + content=( + f"工具 {tool_name} 在本轮中已用相同参数调用过一次，请不要重复调用，" + "应基于现有结果直接回答用户。" + ), + tool_call_id=tc["id"], + ) + ) + trace_entries.append({ + "node": "tools", + "timestamp": time.time(), + "tool_name": tool_name, + "tool_args": tool_args, + "result_status": "duplicate_skipped", + "latency_ms": 0, + }) + continue + attempted_tool_calls[tool_signature] = attempted_tool_calls.get(tool_signature, 0) + 1 + + # Execute via the configured production tool executor. + tool_started_at = time.perf_counter() + result = await tool_executor.execute( + tool_name=tool_name, + tool_args=tool_args, + session_id=session_id, + user_id=user_id, + turn_id=turn_id, + tool_call_id=tc["id"], + agent_id="single-agent", + agent_role="single_agent", + actor_role=actor_role, + ) + measured_tool_latency_ms = max(1, int((time.perf_counter() - tool_started_at) * 1000)) + if not isinstance(result, dict): + result = {"status": "error", "message": str(result)} + if not result.get("_latency_ms"): + result["_latency_ms"] = measured_tool_latency_ms + + # Record trace + trace_entries.append({ + "node": "tools", + "timestamp": time.time(), + "tool_name": tool_name, + "tool_args": tool_args, + "result_status": result.get("status", "unknown"), + "latency_ms": result.get("_latency_ms", 0), + }) + + # Collect ALL confirmation-required results, tagged with tool_call_id + # for correct matching in the orchestrator. + if result.get("status") == "confirmation_required": + pending_confirmations.append({**result, "_tool_call_id": tc["id"]}) + tool_messages.append( + ToolMessage( + content=json.dumps(result, ensure_ascii=False), + tool_call_id=tc["id"], + ) + ) + else: + observation = await _build_tool_observation( + tool_name, result, llm=llm, user_question=_user_question, + ) + tool_messages.append( + ToolMessage(content=observation, tool_call_id=tc["id"]) + ) + + executed = [e.get("tool_name", "?") for e in trace_entries if e.get("node") == "tools"] + if executed: + logger.info("tools_node: executed %s, total_count=%d", executed, new_tool_call_count) + + return { + "messages": tool_messages, + "tool_call_count": new_tool_call_count, + "attempted_tool_calls": attempted_tool_calls, + "pending_confirmations": pending_confirmations, + "trace": trace_entries, + } + + # ----- Node: summarize (forced synthesis when tool limit reached) ----- + async def summarize_node(state: CraterAgentState) -> dict: + """Force LLM to synthesize existing evidence when tool call limit is reached.""" + messages = list(state["messages"]) + + # Inject a nudge so the LLM knows it must summarize now + messages.append( + HumanMessage( + content=( + "[系统提示] 你已达到本轮工具调用上限，无法继续调用工具。" + "请基于已收集到的所有工具返回结果，直接给出完整的综合分析回答。" + ) + ) + ) + + try: + # No tools bound — LLM can only produce text + response = await llm.ainvoke(messages) + except BadRequestError as exc: + if not _is_context_limit_error(exc): + raise + from crater_agent.agent.compaction import compact_messages_with_llm + + compact_messages = await compact_messages_with_llm( + messages, llm, current_query=_extract_current_query(messages), + ) + + if compact_messages is None: + compact_messages = _compact_messages_for_retry(messages) + logger.warning( + "Summarize node context limit hit, retrying with compacted messages: %d -> %d", + len(messages), + len(compact_messages), + ) + response = await llm.ainvoke(compact_messages) + + # Handle qwen thinking mode + if not response.content and not response.tool_calls: + reasoning = getattr(response, "reasoning_content", "") or ( + response.additional_kwargs or {} + ).get("reasoning_content", "") + if reasoning: + response = AIMessage( + content=reasoning, + additional_kwargs=response.additional_kwargs, + response_metadata=response.response_metadata, + ) + + trace_entry = { + "node": "summarize", + "timestamp": time.time(), + "response_length": len(response.content) if response.content else 0, + } + + return { + "messages": [response], + "trace": [trace_entry], + } + + # ----- Conditional edge: should the agent continue? ----- + def should_continue(state: CraterAgentState) -> str: + """Determine if the agent should call tools, wait for confirmation, or respond.""" + messages = state["messages"] + last_message = messages[-1] + tool_call_count = state.get("tool_call_count", 0) + + # Safety: max tool calls reached + if tool_call_count >= settings.max_tool_calls_per_turn: + # If the LLM still wanted to call tools, route to summarize node + # so it can synthesize existing evidence instead of a raw fallback. + if isinstance(last_message, AIMessage) and last_message.tool_calls: + return "summarize" + return "respond" + + # If there are pending confirmations → pause and respond + if state.get("pending_confirmations"): + return "respond" + + # If the LLM produced tool_calls → execute them + if isinstance(last_message, AIMessage) and last_message.tool_calls: + return "tools" + + # Otherwise, LLM decided to respond directly → end + return "respond" + + def after_tools(state: CraterAgentState) -> str: + if state.get("pending_confirmations"): + return "respond" + return "agent" + + # ----- Build the graph ----- + graph = StateGraph(CraterAgentState) + graph.add_node("agent", agent_node) + graph.add_node("tools", tools_node) + graph.add_node("summarize", summarize_node) + + graph.set_entry_point("agent") + graph.add_conditional_edges( + "agent", + should_continue, + { + "tools": "tools", + "summarize": "summarize", + "respond": END, + }, + ) + graph.add_conditional_edges( + "tools", + after_tools, + { + "agent": "agent", + "respond": END, + }, + ) + # summarize always ends the graph + graph.add_edge("summarize", END) + + return graph.compile() diff --git a/crater-agent/crater_agent/agent/prompts.py b/crater-agent/crater_agent/agent/prompts.py new file mode 100644 index 000000000..d57fe293a --- /dev/null +++ b/crater-agent/crater_agent/agent/prompts.py @@ -0,0 +1,75 @@ +"""System prompt templates for the minimal Crater Agent chat.""" + +from __future__ import annotations + +import json +from typing import Any + + +BASE_PROMPT = """\ +你是 Crater AI Agent，负责帮助用户查询、诊断和管理自己的平台作业。 + +工作原则： +1. 使用中文回答，结构尽量是：结论 -> 证据 -> 建议。 +2. 用户询问作业状态、失败原因、日志、事件、指标或排队原因时，必须先调用相关只读工具取证。 +3. `job_name` 是平台作业系统名，例如 `jpt-xxx`、`sg-xxx`；不要把显示名当成 job_name。 +4. 停止、删除、重提、创建作业必须调用对应写工具触发系统确认卡片，不要在文字里伪造确认表单。 + - 如果用户本轮已经表达了删除/停止/重提/创建意图，且你已经能确定目标 job_name 或创建参数，必须立刻调用对应写工具。 + - 不要先回复“是否确认删除/是否确认创建/请确认后我再执行”；确认由系统卡片完成。 + - 即使目标是通过 list_user_jobs 等只读工具刚刚匹配出来的，也应在同一轮继续调用写工具。 +5. 如果工具返回无权限、403、forbidden 或 not found，要直接说明“该对象不存在或你没有访问权限”，不要说成临时故障。 +6. 信息不足时先澄清；不要编造作业状态、日志、事件、节点或配额。 + +可用只读工具： +- get_job_detail / get_job_events / get_job_logs / diagnose_job / get_diagnostic_context +- search_similar_failures / query_job_metrics / analyze_queue_status +- get_realtime_capacity / check_quota / list_user_jobs +- list_available_images / list_available_gpu_models / get_job_templates / get_resource_recommendation + +可用写工具： +- create_jupyter_job / create_webide_job / create_custom_job / create_pytorch_job / create_tensorflow_job +- resubmit_job / stop_job / delete_job +""" + +FIRST_TIME_ADDON = """\ + +首次使用提示： +- 可以问“我的作业为什么失败了” +- 可以问“帮我看 job_name 的日志” +- 可以让我创建、停止、删除或重新提交作业，系统会先弹确认卡片 +""" + + +def _json_block(value: Any) -> str: + try: + return json.dumps(value or {}, ensure_ascii=False, indent=2) + except TypeError: + return "{}" + + +def build_system_prompt( + context: dict, + is_first_time: bool = False, + user_message: str = "", +) -> str: + actor = context.get("actor") or {} + page = context.get("page") or {} + capabilities = context.get("capabilities") or {} + prompt = BASE_PROMPT + if is_first_time: + prompt += FIRST_TIME_ADDON + prompt += f""" + +当前上下文： +- 用户: {actor.get("username") or actor.get("user_name") or ""} +- 用户 ID: {actor.get("user_id") or ""} +- 账户 ID: {actor.get("account_id") or ""} +- 页面: {page.get("url") or page.get("route") or ""} +- 本轮用户输入: {user_message} + +工具能力摘要： +```json +{_json_block(capabilities)} +``` +""" + return prompt diff --git a/crater-agent/crater_agent/agent/state.py b/crater-agent/crater_agent/agent/state.py new file mode 100644 index 000000000..d1c85a27a --- /dev/null +++ b/crater-agent/crater_agent/agent/state.py @@ -0,0 +1,38 @@ +"""LangGraph State definition for Crater Agent.""" + +from __future__ import annotations + +import operator +from typing import Annotated, Any + +from langgraph.graph import MessagesState + + +class CraterAgentState(MessagesState): + """State that persists through the entire ReAct loop. + + Inherits `messages` from MessagesState (list of BaseMessage with add_messages reducer). + """ + + # User and page context injected by Go backend + context: dict[str, Any] + + # Track tool call count to prevent infinite loops + tool_call_count: int + + # Track same-turn tool invocations to avoid repeated calls with identical args + attempted_tool_calls: dict[str, int] + + # When write operations need user confirmation, store them here. + # Each entry is the raw confirmation result dict from Go backend, + # tagged with "_tool_call_id" for matching to the LLM's tool_call. + # The ReAct loop pauses and returns these to the frontend. + pending_confirmations: list[dict[str, Any]] + + # When True, the tool call limit was reached and the agent should + # summarize existing evidence instead of calling more tools. + force_limit_reached: bool + + # Accumulated trace records for evaluation/auditing + # Uses operator.add reducer so each node's trace entries are appended, not replaced + trace: Annotated[list[dict[str, Any]], operator.add] diff --git a/crater-agent/crater_agent/app.py b/crater-agent/crater_agent/app.py new file mode 100644 index 000000000..fd38dab25 --- /dev/null +++ b/crater-agent/crater_agent/app.py @@ -0,0 +1,165 @@ +"""FastAPI application for Crater Agent Service. + +Exposes: +- POST /chat — accepts user message + context, returns SSE stream of agent events +- GET /health — health check +""" + +from __future__ import annotations + +import json +import logging +from typing import Any + +# Configure logging before anything else — without this, logger.info() in agent +# modules is silently swallowed because the root logger defaults to WARNING. +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s [%(levelname)s] %(name)s: %(message)s", + datefmt="%H:%M:%S", +) + +from fastapi import FastAPI + +logger = logging.getLogger(__name__) +from pydantic import BaseModel, Field +from sse_starlette.sse import EventSourceResponse + +from crater_agent.config import settings +from crater_agent.llm.client import ( + ModelClientFactory, + local_llm_config_fallback_enabled, + normalize_runtime_llm_client_configs, + reset_runtime_llm_client_configs, + set_runtime_llm_client_configs, +) +from crater_agent.orchestrators.single import SingleAgentOrchestrator + +app = FastAPI(title="Crater Agent Service", version="0.1.0") + +logger.info( + "agent service starting: backend=%s", + settings.crater_backend_url, +) + +single_orchestrator = SingleAgentOrchestrator() + + +class ChatRequest(BaseModel): + """Request from Go backend.""" + + session_id: str + message: str + turn_id: str | None = None + context: dict[str, Any] = Field(default_factory=dict) + user_id: int | None = None + account_id: int | None = None + username: str | None = None + page_context: dict[str, Any] | None = None + + +def build_request_context(request: ChatRequest) -> dict[str, Any]: + context = dict(request.context or {}) + actor = dict(context.get("actor") or {}) + if request.user_id is not None and "user_id" not in actor: + actor["user_id"] = request.user_id + if request.account_id is not None and "account_id" not in actor: + actor["account_id"] = request.account_id + if request.username is not None and "username" not in actor: + actor["username"] = request.username + if actor: + context["actor"] = actor + if request.page_context and "page" not in context: + context["page"] = request.page_context + return context + + +def get_orchestration_mode(context: dict[str, Any]) -> str: + return "single_agent" + + +def get_request_model_factory(context: dict[str, Any]) -> ModelClientFactory: + runtime_clients = get_request_llm_client_configs(context) + if runtime_clients: + return ModelClientFactory(raw_clients=runtime_clients) + return ModelClientFactory() + + +def get_request_llm_client_configs(context: dict[str, Any]) -> dict[str, Any] | None: + llm_context = context.get("llm") if isinstance(context, dict) else {} + if isinstance(llm_context, dict): + raw_clients = llm_context.get("client_config") + return normalize_runtime_llm_client_configs(raw_clients) + return None + + +@app.get("/health") +async def health(): + default_model = "" + if local_llm_config_fallback_enabled(): + try: + default_model = str(settings.get_llm_client_config("default").get("model") or "") + except Exception: + pass + try: + default_model = ModelClientFactory().create("default").model_name + except Exception: + pass + return {"status": "ok", "model": default_model} + + +@app.get("/config-summary") +async def config_summary(): + if not local_llm_config_fallback_enabled(): + return { + "defaultOrchestrationMode": settings.normalized_default_orchestration_mode(), + "availableModes": ["single_agent"], + "localLLMConfigFallbackEnabled": False, + } + return settings.public_agent_config_summary() + + +@app.post("/chat") +async def chat(request: ChatRequest): + """Process a chat message and return SSE stream. + + The Go backend calls this endpoint with the user message and context. + Returns Server-Sent Events with thinking, tool_call, tool_result, message events. + """ + async def event_generator(): + request_context = build_request_context(request) + request.context = request_context + orchestration_mode = get_orchestration_mode(request_context) + runtime_llm_clients = get_request_llm_client_configs(request_context) + + runtime_token = set_runtime_llm_client_configs(runtime_llm_clients) + try: + model_factory = get_request_model_factory(request_context) + orchestrator = single_orchestrator + async for event in orchestrator.stream(request=request, model_factory=model_factory): + yield { + "event": event["event"], + "data": json.dumps(event.get("data", {}), ensure_ascii=False), + } + except Exception as e: + logger.exception("Orchestrator stream error") + yield { + "event": "error", + "data": json.dumps( + {"code": "agent_error", "message": str(e)}, ensure_ascii=False + ), + } + yield { + "event": "done", + "data": json.dumps({}, ensure_ascii=False), + } + finally: + reset_runtime_llm_client_configs(runtime_token) + + return EventSourceResponse(event_generator(), ping=0) + + +if __name__ == "__main__": + import uvicorn + + uvicorn.run(app, host=settings.host, port=settings.port) diff --git a/crater-agent/crater_agent/config.py b/crater-agent/crater_agent/config.py new file mode 100644 index 000000000..5baeadae6 --- /dev/null +++ b/crater-agent/crater_agent/config.py @@ -0,0 +1,151 @@ +"""Configuration for Crater Agent Service.""" + +import json +from pathlib import Path +from typing import Any + +from pydantic import Field +from pydantic_settings import BaseSettings + + +class Settings(BaseSettings): + """Agent service configuration, loaded from environment variables.""" + + default_orchestration_mode: str = Field( + default="single_agent", + description="Default orchestration mode for agent chat", + ) + + # Single source of truth for LLM routing. + # This file contains the direct purpose/role -> client config map. + llm_clients_config_path: str = Field( + default="./config/llm-clients.json", + description="Path to the LLM client map JSON file", + ) + + backend_debug_config_path: str = Field( + default="", + description="Optional local-dev path to backend debug config YAML for platform discovery", + ) + + # Crater Go Backend + crater_backend_url: str = Field( + default="http://localhost:8080", description="Crater Go backend URL" + ) + crater_backend_internal_token: str = Field( + default="", description="Shared token for Python Agent -> Go internal tool execution" + ) + agent_internal_token: str = Field( + default="dev-agent-internal-token", + description="Internal token for Go backend authentication", + ) + + # Agent Behavior + max_tool_calls_per_turn: int = Field( + default=15, description="Max tool calls in a single ReAct loop" + ) + tool_execution_timeout: int = Field(default=30, description="Tool execution timeout (seconds)") + history_max_tokens: int = Field( + default=4000, description="Max tokens for conversation history" + ) + max_context_tokens: int = Field( + default=30000, description="Estimated LLM context window budget for proactive compaction" + ) + tokenizer_encoding: str = Field( + default="cl100k_base", description="tiktoken encoding name for token counting" + ) + + # Service + host: str = Field(default="0.0.0.0") + port: int = Field(default=8000) + debug: bool = Field(default=False) + + model_config = { + "env_prefix": "CRATER_AGENT_", + "env_file": ".env", + # .env may also contain non-prefixed secret variables such as DASHSCOPE_API_KEY + # referenced indirectly by llm client configs. They should not fail Settings init. + "extra": "ignore", + } + + def normalized_default_orchestration_mode(self) -> str: + return "single_agent" + + def resolve_llm_clients_config_path(self) -> Path: + configured = self.llm_clients_config_path.strip() or "./config/llm-clients.json" + return self._resolve_config_path(configured) + + def resolve_backend_debug_config_path(self) -> Path | None: + configured = self.backend_debug_config_path.strip() + if not configured: + return None + return self._resolve_config_path(configured) + + def _resolve_config_path(self, configured: str) -> Path: + raw_path = Path(configured).expanduser() + if raw_path.is_absolute(): + return raw_path + + cwd_candidate = Path.cwd() / raw_path + if cwd_candidate.exists(): + return cwd_candidate + + project_root = Path(__file__).resolve().parents[1] + return project_root / raw_path + + def load_llm_client_configs(self) -> dict[str, dict[str, Any]]: + path = self.resolve_llm_clients_config_path() + if not path.exists(): + raise ValueError(f"LLM client config file not found: {path}") + + raw = path.read_text(encoding="utf-8").strip() + if not raw: + raise ValueError(f"LLM client config file is empty: {path}") + + loaded = json.loads(raw) + if not isinstance(loaded, dict): + raise ValueError("LLM client config must decode to a JSON object") + + configs = { + str(name): dict(config) + for name, config in loaded.items() + if isinstance(name, str) and isinstance(config, dict) + } + if "default" not in configs: + raise ValueError("LLM client config must define a 'default' client") + return configs + + def get_llm_client_config(self, client_key: str = "default") -> dict[str, Any]: + configs = self.load_llm_client_configs() + normalized = client_key.strip() + if normalized and normalized in configs: + return dict(configs[normalized]) + return dict(configs["default"]) + + def public_agent_config_summary(self) -> dict[str, Any]: + configs = self.load_llm_client_configs() + return { + "defaultOrchestrationMode": self.normalized_default_orchestration_mode(), + "availableModes": ["single_agent"], + "llmConfigPath": str(self.resolve_llm_clients_config_path()), + "backendDebugConfigPath": ( + str(self.resolve_backend_debug_config_path()) + if self.resolve_backend_debug_config_path() is not None + else "" + ), + "llmClientKeys": list(configs.keys()), + "llmClients": { + name: { + "baseUrl": str(config.get("base_url") or ""), + "model": str(config.get("model") or ""), + "temperature": float(config.get("temperature") or 0.0), + "maxTokens": int(config.get("max_tokens") or 0), + "timeout": int(config.get("timeout") or 0), + "apiKeyEnv": str(config.get("api_key_env") or ""), + "hasInlineApiKey": bool(str(config.get("api_key") or "").strip()), + } + for name, config in configs.items() + }, + } + +settings = Settings() diff --git a/crater-agent/crater_agent/llm/__init__.py b/crater-agent/crater_agent/llm/__init__.py new file mode 100644 index 000000000..678d5489d --- /dev/null +++ b/crater-agent/crater_agent/llm/__init__.py @@ -0,0 +1 @@ +"""LLM helpers for Crater Agent.""" diff --git a/crater-agent/crater_agent/llm/client.py b/crater-agent/crater_agent/llm/client.py new file mode 100644 index 000000000..9acb9bb9d --- /dev/null +++ b/crater-agent/crater_agent/llm/client.py @@ -0,0 +1,196 @@ +"""LLM client factory for Crater Agent.""" + +from __future__ import annotations + +import os +from contextvars import ContextVar +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import httpx +from dotenv import dotenv_values +from langchain_openai import ChatOpenAI + +from crater_agent.config import settings + +NO_AUTH_API_KEY_PLACEHOLDER = "sk-no-auth-required" + +# Load .env once at module level so api_key_env references can resolve +_dotenv_cache: dict[str, str | None] | None = None +_runtime_llm_clients: ContextVar[dict[str, Any] | None] = ContextVar( + "runtime_llm_clients", + default=None, +) + + +def get_runtime_llm_client_configs() -> dict[str, Any] | None: + return _runtime_llm_clients.get() + + +def set_runtime_llm_client_configs(raw_clients: dict[str, Any] | None): + normalized = normalize_runtime_llm_client_configs(raw_clients) + return _runtime_llm_clients.set(normalized) + + +def reset_runtime_llm_client_configs(token) -> None: + _runtime_llm_clients.reset(token) + + +def normalize_runtime_llm_client_configs(raw_clients: dict[str, Any] | None) -> dict[str, Any] | None: + if not isinstance(raw_clients, dict): + return None + default = raw_clients.get("default") + if not isinstance(default, dict): + return None + return { + str(name): dict(config) + for name, config in raw_clients.items() + if isinstance(name, str) and isinstance(config, dict) + } + + +def local_llm_config_fallback_enabled() -> bool: + """Return whether the agent may use its local llm-clients.json as a fallback.""" + + value = os.getenv("CRATER_AGENT_ALLOW_LOCAL_LLM_CONFIG_FALLBACK", "").strip().lower() + return value in {"1", "true", "yes", "on"} + + +def _load_dotenv() -> dict[str, str | None]: + global _dotenv_cache + if _dotenv_cache is None: + env_path = Path(settings.llm_clients_config_path).parent.parent / ".env" + if not env_path.exists(): + env_path = Path.cwd() / ".env" + _dotenv_cache = dotenv_values(env_path) if env_path.exists() else {} + return _dotenv_cache + + +@dataclass +class ClientConfig: + """Direct LLM client config for one role or purpose.""" + + name: str + provider: str = "openai_compatible" + base_url: str = "" + api_key: str = "" + api_key_env: str = "" + model: str = "" + temperature: float = 0.1 + max_tokens: int = 4096 + timeout: int = 60 + max_retries: int = 2 + verify_ssl: bool = True + streaming: bool = False + stream_usage: bool | None = None + headers: dict[str, str] = field(default_factory=dict) + model_kwargs: dict[str, Any] = field(default_factory=dict) + + @classmethod + def from_dict(cls, name: str, raw: dict[str, Any]) -> "ClientConfig": + return cls( + name=name, + provider=str(raw.get("provider") or "openai_compatible"), + base_url=str(raw.get("base_url") or ""), + api_key=str(raw.get("api_key") or ""), + api_key_env=str(raw.get("api_key_env") or ""), + model=str(raw.get("model") or ""), + temperature=float(raw.get("temperature") or 0.1), + max_tokens=int(raw.get("max_tokens") or 4096), + timeout=int(raw.get("timeout") or 60), + max_retries=int(raw.get("max_retries") or 2), + verify_ssl=bool(raw.get("verify_ssl", True)), + streaming=bool(raw.get("streaming", False)), + stream_usage=( + bool(raw.get("stream_usage")) + if raw.get("stream_usage") is not None + else None + ), + headers={str(key): str(value) for key, value in (raw.get("headers") or {}).items()}, + model_kwargs={str(key): value for key, value in (raw.get("model_kwargs") or {}).items()}, + ) + + def resolved_api_key(self) -> str: + if self.api_key_env: + # 1. Check process environment + env_value = os.getenv(self.api_key_env, "").strip() + if env_value: + return env_value + # 2. Check .env file (Pydantic only loads CRATER_AGENT_ prefixed vars) + dotenv = _load_dotenv() + dotenv_value = str(dotenv.get(self.api_key_env) or "").strip() + if dotenv_value: + return dotenv_value + if self.api_key.strip(): + return self.api_key.strip() + return NO_AUTH_API_KEY_PLACEHOLDER + + +def normalize_client_map(raw: dict[str, Any]) -> dict[str, ClientConfig]: + """Normalize the direct purpose/role -> client config map.""" + + lookup: dict[str, ClientConfig] = {} + for name, cfg in raw.items(): + if isinstance(cfg, dict): + lookup[str(name)] = ClientConfig.from_dict(str(name), cfg) + + fallback = lookup.get("default") + if fallback is None: + raise ValueError("LLM client config must define a 'default' client") + + for name, cfg in list(lookup.items()): + if not cfg.base_url: + cfg.base_url = fallback.base_url + if not cfg.model: + cfg.model = fallback.model + if not cfg.max_tokens: + cfg.max_tokens = fallback.max_tokens + if not cfg.timeout: + cfg.timeout = fallback.timeout + lookup[name] = cfg + + return lookup + + +class ModelClientFactory: + """Builds ChatOpenAI clients from a direct client-key -> config map.""" + + def __init__(self, raw_clients: dict[str, Any] | None = None): + if raw_clients is None: + raw_clients = get_runtime_llm_client_configs() + if raw_clients is None and local_llm_config_fallback_enabled(): + raw_clients = settings.load_llm_client_configs() + if raw_clients is None: + raise ValueError( + "Platform LLM config is required; local llm-clients.json fallback is disabled" + ) + self._clients = normalize_client_map(raw_clients) + + @property + def client_map(self) -> dict[str, ClientConfig]: + return self._clients + + def create(self, client_key: str = "default") -> ChatOpenAI: + normalized_key = str(client_key or "").strip() or "default" + config = self._clients.get(normalized_key) or self._clients["default"] + client_kwargs: dict[str, Any] = { + "base_url": config.base_url, + "api_key": config.resolved_api_key(), + "model": config.model, + "temperature": config.temperature, + "max_tokens": config.max_tokens, + "timeout": config.timeout, + "max_retries": config.max_retries, + "streaming": config.streaming, + } + if config.stream_usage is not None: + client_kwargs["stream_usage"] = config.stream_usage + if not config.verify_ssl: + client_kwargs["http_client"] = httpx.Client(verify=False) + client_kwargs["http_async_client"] = httpx.AsyncClient(verify=False) + if config.headers: + client_kwargs["default_headers"] = config.headers + if config.model_kwargs: + client_kwargs["model_kwargs"] = config.model_kwargs + return ChatOpenAI(**client_kwargs) diff --git a/crater-agent/crater_agent/llm/tokenizer.py b/crater-agent/crater_agent/llm/tokenizer.py new file mode 100644 index 000000000..bd6542f9f --- /dev/null +++ b/crater-agent/crater_agent/llm/tokenizer.py @@ -0,0 +1,77 @@ +"""Token counting with tiktoken backend and heuristic fallback. + +Provides accurate token estimation for context window management. +Uses tiktoken cl100k_base encoding (compatible with Qwen/GPT-4 family). +Falls back to character-based heuristic if tiktoken is unavailable. + +All counting is local — zero network requests, zero token consumption. +""" + +from __future__ import annotations + +import logging +from typing import Any + +logger = logging.getLogger(__name__) + +_MESSAGE_OVERHEAD_TOKENS = 4 # per-message framing (role, separators) + + +class TokenCounter: + """Token counter with tiktoken backend and heuristic fallback.""" + + def __init__(self, encoding_name: str = "cl100k_base"): + self._encoding = None + try: + import tiktoken + + self._encoding = tiktoken.get_encoding(encoding_name) + except Exception: + logger.info("tiktoken unavailable, using heuristic token estimation") + + def count_text(self, text: str) -> int: + """Count tokens in a text string.""" + if not text: + return 0 + if self._encoding is not None: + return len(self._encoding.encode(text)) + return _heuristic_count(text) + + def count_messages(self, messages: list[Any]) -> int: + """Count tokens across a list of LangChain messages.""" + total = 0 + for msg in messages: + content = str(getattr(msg, "content", "") or "") + total += self.count_text(content) + _MESSAGE_OVERHEAD_TOKENS + return total + + +def _heuristic_count(text: str) -> int: + """Fallback heuristic: ~1 token per 2 CJK chars, ~1 token per 4 Latin chars.""" + cjk = sum(1 for c in text if "\u4e00" <= c <= "\u9fff") + latin = len(text) - cjk + return cjk // 2 + latin // 4 + 1 + + +# Module-level singleton, lazily initialized. +_default_counter: TokenCounter | None = None + + +def get_token_counter() -> TokenCounter: + """Get the module-level TokenCounter singleton.""" + global _default_counter + if _default_counter is None: + from crater_agent.config import settings + + _default_counter = TokenCounter(encoding_name=settings.tokenizer_encoding) + return _default_counter + + +def count_tokens(text: str) -> int: + """Convenience function: count tokens in text using the default counter.""" + return get_token_counter().count_text(text) + + +def count_message_tokens(messages: list[Any]) -> int: + """Convenience function: count tokens across messages using the default counter.""" + return get_token_counter().count_messages(messages) diff --git a/crater-agent/crater_agent/memory/__init__.py b/crater-agent/crater_agent/memory/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/crater-agent/crater_agent/memory/session.py b/crater-agent/crater_agent/memory/session.py new file mode 100644 index 000000000..0c67cd6bf --- /dev/null +++ b/crater-agent/crater_agent/memory/session.py @@ -0,0 +1,84 @@ +"""Session memory management. + +Handles conversation history loading and token budget management. +History is provided by the Go backend (loaded from PostgreSQL) and +passed in the request context. +""" + +from __future__ import annotations + +from langchain_core.messages import AIMessage, HumanMessage + +from crater_agent.llm.tokenizer import count_tokens + + +def estimate_tokens(text: str) -> int: + """Count tokens in text using tiktoken (with heuristic fallback).""" + return count_tokens(text) + + +def _truncate_head_tail(text: str, max_chars: int) -> str: + """Truncate keeping head and tail for better context preservation.""" + if len(text) <= max_chars: + return text + half = max_chars // 2 + return f"{text[:half]}\n\n...(内容过长，已截断)...\n\n{text[-half:]}" + + +def build_history_messages( + history: list[dict], + max_tokens: int = 4000, + tool_result_max_chars: int = 1200, + tool_error_max_chars: int = 1600, +) -> list: + """Build LangChain message objects from Go-provided history. + + Loads messages from most recent backwards until token budget is exhausted. + Tool results are truncated (head+tail) to save tokens while preserving + key information at both ends of the output. + + Args: + history: List of message dicts from Go backend + [{"role": "user", "content": "..."}, ...] + max_tokens: Maximum token budget for history + tool_result_max_chars: Max chars for tool result content + tool_error_max_chars: Max chars for tool error content + + Returns: + List of LangChain message objects, in chronological order + """ + if not history: + return [] + + selected = [] + token_count = 0 + + for msg in reversed(history): + role = msg.get("role", "") + content = msg.get("content", "") + + # Truncate tool results with head+tail to preserve key info + if role == "tool": + is_error = any(kw in content for kw in ("error", "Error", "failed", "Failed", "错误", "失败")) + limit = tool_error_max_chars if is_error else tool_result_max_chars + content = _truncate_head_tail(content, limit) + + msg_tokens = estimate_tokens(content) + if token_count + msg_tokens > max_tokens: + break + + if role == "user": + selected.append(HumanMessage(content=content)) + elif role == "assistant": + selected.append(AIMessage(content=content)) + elif role == "tool": + tool_call_id = str(msg.get("tool_call_id", "unknown") or "unknown").strip() + selected.append( + AIMessage( + content=f"【历史工具结果 {tool_call_id}】{content}", + ) + ) + + token_count += msg_tokens + + return list(reversed(selected)) diff --git a/crater-agent/crater_agent/orchestrators/__init__.py b/crater-agent/crater_agent/orchestrators/__init__.py new file mode 100644 index 000000000..6b5003ac6 --- /dev/null +++ b/crater-agent/crater_agent/orchestrators/__init__.py @@ -0,0 +1 @@ +"""Agent orchestrators.""" diff --git a/crater-agent/crater_agent/orchestrators/single.py b/crater-agent/crater_agent/orchestrators/single.py new file mode 100644 index 000000000..09b69c464 --- /dev/null +++ b/crater-agent/crater_agent/orchestrators/single.py @@ -0,0 +1,512 @@ +"""Single-agent orchestrator wrapper.""" + +from __future__ import annotations + +import json +import re +import time +from typing import Any, AsyncIterator + +from langchain_core.messages import AIMessage, HumanMessage + +from crater_agent.agent.graph import create_agent_graph +from crater_agent.config import settings +from crater_agent.llm.client import ModelClientFactory +from crater_agent.memory.session import build_history_messages +from crater_agent.tools.executor import CompositeToolExecutor, ToolExecutorProtocol +from crater_agent.tools.tool_selector import sanitize_capabilities_for_context + + +def _extract_llm_usage(output: Any) -> dict[str, int]: + usage = getattr(output, "usage_metadata", None) or {} + response_metadata = getattr(output, "response_metadata", None) or {} + token_usage = ( + response_metadata.get("token_usage") if isinstance(response_metadata, dict) else {} + ) or {} + input_tokens = ( + usage.get("input_tokens") + or usage.get("prompt_tokens") + or token_usage.get("prompt_tokens") + or token_usage.get("input_tokens") + or 0 + ) + output_tokens = ( + usage.get("output_tokens") + or usage.get("completion_tokens") + or token_usage.get("completion_tokens") + or token_usage.get("output_tokens") + or 0 + ) + total_tokens = ( + usage.get("total_tokens") + or token_usage.get("total_tokens") + or (int(input_tokens or 0) + int(output_tokens or 0)) + ) + has_reported_usage = bool(usage) or bool(token_usage) + return { + "llm_input_tokens": int(input_tokens or 0), + "llm_output_tokens": int(output_tokens or 0), + "total_tokens": int(total_tokens or 0), + "llm_reported_token_calls": 1 if has_reported_usage else 0, + "llm_missing_token_calls": 0 if has_reported_usage else 1, + } + + +def _looks_like_continuation_reply(user_message: str) -> bool: + normalized = str(user_message or "").strip().lower() + if not normalized: + return False + if normalized in {"确认", "继续", "这个", "就这个", "好的", "ok", "yes", "1", "2", "3"}: + return True + if re.fullmatch(r"[A-Za-z0-9]+(?:-[A-Za-z0-9]+){2,}", normalized): + return True + return any( + token in normalized + for token in ( + "第一个", + "第1个", + "第二个", + "第2个", + "改成", + "名字叫", + "重新来", + "继续刚才", + "按刚才", + "全部", + "所有", + ) + ) + + +class SingleAgentOrchestrator: + def __init__(self, tool_executor: ToolExecutorProtocol | None = None): + self.tool_executor = tool_executor or CompositeToolExecutor() + + async def stream(self, *, request: Any, model_factory: ModelClientFactory) -> AsyncIterator[dict]: + llm = model_factory.create("default") + graph = create_agent_graph(tool_executor=self.tool_executor, llm=llm) + context = dict(request.context or {}) + context["capabilities"] = sanitize_capabilities_for_context(context, context.get("capabilities")) + pending_tool_calls: list[dict[str, Any]] = [] + tool_result_summaries: list[str] = [] + pending_final_content = "" + emitted_final_answer = False + emitted_confirmation = False + llm_started_at: float | None = None + usage_summary: dict[str, int] = { + "llm_calls": 0, + "llm_input_tokens": 0, + "llm_output_tokens": 0, + "total_tokens": 0, + "llm_reported_token_calls": 0, + "llm_missing_token_calls": 0, + "reported_token_coverage": 0, + "llm_latency_ms": 0, + "tool_latency_ms": 0, + "tool_calls": 0, + "read_tool_calls": 0, + "write_tool_calls": 0, + "evidence_items": 0, + } + initial_state = { + "messages": [HumanMessage(content=request.message)], + "context": { + **context, + "session_id": request.session_id, + "turn_id": request.turn_id, + }, + "tool_call_count": 0, + "attempted_tool_calls": {}, + "pending_confirmations": [], + "trace": [], + } + history = context.get("history", []) + if history: + initial_state["messages"] = build_history_messages( + history=history, + max_tokens=settings.history_max_tokens, + tool_result_max_chars=160, + ) + initial_state["messages"] + continuation = context.get("continuation") or {} + current_request_is_follow_up = _looks_like_continuation_reply(request.message) + continuation_messages: list[HumanMessage] = [] + pending_confirmation = continuation.get("pending_confirmation") or {} + if isinstance(pending_confirmation, dict) and pending_confirmation and current_request_is_follow_up: + pending_summary = { + "tool_name": pending_confirmation.get("tool_name", ""), + "result_status": pending_confirmation.get("result_status", ""), + "tool_args": pending_confirmation.get("tool_args", {}), + } + continuation_messages.append( + HumanMessage( + content=( + "[系统续接上下文] 上一轮仍有待确认操作：" + f"{json.dumps(pending_summary, ensure_ascii=False)}。" + "只有当用户当前输入明显是在继续这件事时，才沿用该上下文；否则按新请求处理。" + ) + ) + ) + elif isinstance(pending_confirmation, dict) and pending_confirmation: + continuation_messages.append( + HumanMessage( + content=( + "[系统续接上下文] 上一轮仍有待确认操作，但当前用户输入看起来是新的独立请求。" + "除非用户明确说“确认/继续/这个/第一个/具体名称”，否则不要延续上一轮创建或修改计划；" + "优先回答本轮问题。" + ) + ) + ) + resume_after_confirmation = continuation.get("resume_after_confirmation") or {} + if isinstance(resume_after_confirmation, dict) and resume_after_confirmation and current_request_is_follow_up: + resume_summary = { + "tool_name": resume_after_confirmation.get("tool_name", ""), + "result_status": resume_after_confirmation.get("result_status", ""), + "confirmed": resume_after_confirmation.get("confirmed"), + "tool_args": resume_after_confirmation.get("tool_args", {}), + "result": resume_after_confirmation.get("result", {}), + } + continuation_messages.append( + HumanMessage( + content=( + "[系统续接上下文] 上一轮确认结果：" + f"{json.dumps(resume_summary, ensure_ascii=False)}。" + "当前轮应基于这个结果理解用户后续输入，但不要忽略本轮新的完整请求。" + ) + ) + ) + elif isinstance(resume_after_confirmation, dict) and resume_after_confirmation: + continuation_messages.append( + HumanMessage( + content=( + "[系统续接上下文] 上一轮确认流程已经结束，但当前用户输入看起来不是续接语句。" + "不要把上一轮写操作当作当前默认目标；若本轮是在问失败原因、状态或教程，" + "就按新的诊断/检索请求处理。" + ) + ) + ) + if continuation_messages: + initial_state["messages"] = ( + initial_state["messages"][:-1] + + continuation_messages + + initial_state["messages"][-1:] + ) + + yield { + "event": "agent_run_started", + "data": { + "turnId": request.turn_id, + "sessionId": request.session_id, + "agentId": "single-agent", + "agentRole": "single_agent", + "status": "started", + "summary": "单核 Agent 已启动", + }, + } + + async for event in graph.astream_events(initial_state, version="v2"): + kind = event["event"] + if kind == "on_chat_model_start": + llm_started_at = time.monotonic() + yield { + "event": "agent_status", + "data": { + "turnId": request.turn_id, + "agentId": "single-agent", + "agentRole": "single_agent", + "status": "running", + "summary": "Agent 思考中", + }, + } + continue + + if kind == "on_chat_model_end": + output = event["data"]["output"] + latency_ms = ( + int((time.monotonic() - llm_started_at) * 1000) + if llm_started_at is not None + else 0 + ) + usage = _extract_llm_usage(output) + usage_summary["llm_calls"] += 1 + usage_summary["llm_input_tokens"] += usage["llm_input_tokens"] + usage_summary["llm_output_tokens"] += usage["llm_output_tokens"] + usage_summary["total_tokens"] += usage["total_tokens"] + usage_summary["llm_reported_token_calls"] += usage["llm_reported_token_calls"] + usage_summary["llm_missing_token_calls"] += usage["llm_missing_token_calls"] + usage_summary["reported_token_coverage"] = ( + usage_summary["llm_reported_token_calls"] / usage_summary["llm_calls"] + if usage_summary["llm_calls"] + else 0 + ) + usage_summary["llm_latency_ms"] += latency_ms + yield { + "event": "llm_call_completed", + "data": { + "turnId": request.turn_id, + "agentId": "single-agent", + "agentRole": "single_agent", + "latencyMs": latency_ms, + "usage": usage, + }, + } + llm_started_at = None + if isinstance(output, AIMessage): + if output.tool_calls: + for tc in output.tool_calls: + pending_tool_calls.append(tc) + yield { + "event": "tool_call_started", + "data": { + "turnId": request.turn_id, + "agentId": "single-agent", + "agentRole": "single_agent", + "toolCallId": tc.get("id"), + "toolName": tc["name"], + "toolArgs": tc["args"], + "status": "executing", + }, + } + else: + # Extract content — handle qwen thinking mode where content + # may be empty but reasoning_content holds the actual reply. + final_content = output.content or "" + if not final_content: + # qwen3 thinking mode: content in additional_kwargs + final_content = ( + getattr(output, "reasoning_content", "") + or (output.additional_kwargs or {}).get("reasoning_content", "") + ) + if final_content: + pending_final_content = final_content + continue + + if kind == "on_tool_end": + output = event["data"].get("output", "") + raw_output = getattr(output, "content", output) + pending = pending_tool_calls.pop(0) if pending_tool_calls else {} + tool_name = pending.get("name") or event.get("name", "unknown") + tool_call_id = pending.get("id") + + try: + result_data = raw_output if isinstance(raw_output, dict) else json.loads(str(raw_output)) + except (json.JSONDecodeError, TypeError): + result_data = None + + if isinstance(result_data, dict) and result_data.get("status") == "confirmation_required": + confirmation = result_data.get("confirmation", {}) + tool_latency_ms = int(result_data.get("_latency_ms") or 0) + emitted_confirmation = True + usage_summary["tool_calls"] += 1 + usage_summary["write_tool_calls"] += 1 + usage_summary["evidence_items"] += 1 + usage_summary["tool_latency_ms"] += tool_latency_ms + yield { + "event": "tool_call_confirmation_required", + "data": { + "turnId": request.turn_id, + "agentId": "single-agent", + "agentRole": "single_agent", + "toolCallId": tool_call_id, + "confirmId": confirmation.get("confirm_id", ""), + "action": confirmation.get("tool_name", tool_name), + "description": confirmation.get("description", ""), + "riskLevel": confirmation.get("risk_level", ""), + "permissionExplanation": confirmation.get( + "permission_explanation", "" + ), + "riskExplanation": confirmation.get("risk_explanation", ""), + "affectedResources": confirmation.get( + "affected_resources", [] + ), + "interaction": confirmation.get("interaction", "approval"), + "form": confirmation.get("form"), + "status": "awaiting_confirmation", + "latencyMs": tool_latency_ms, + }, + } + continue + + tool_result_summaries.append(f"{tool_name}: {str(raw_output)[:300]}") + tool_latency_ms = ( + int(result_data.get("_latency_ms") or 0) + if isinstance(result_data, dict) + else 0 + ) + usage_summary["tool_calls"] += 1 + usage_summary["read_tool_calls"] += 1 + usage_summary["evidence_items"] += 1 + usage_summary["tool_latency_ms"] += tool_latency_ms + yield { + "event": "tool_call_completed", + "data": { + "turnId": request.turn_id, + "agentId": "single-agent", + "agentRole": "single_agent", + "toolCallId": tool_call_id, + "toolName": tool_name, + "result": raw_output, + "resultSummary": str(raw_output)[:500], + "status": "error" if isinstance(result_data, dict) and result_data.get("status") == "error" else "done", + "isError": isinstance(result_data, dict) and result_data.get("status") == "error", + "latencyMs": tool_latency_ms, + }, + } + continue + + if kind == "on_chain_end" and event.get("name") == "tools": + output = event["data"].get("output", {}) + if not pending_tool_calls: + continue + if not isinstance(output, dict): + continue + + # Build a lookup of confirmation results keyed by tool_call_id + # so we can match each confirmation to the correct pending_tool_call. + pending_confs = output.get("pending_confirmations") or [] + conf_by_tc_id: dict[str, dict] = {} + for conf in pending_confs: + tc_id = conf.get("_tool_call_id") + if tc_id: + conf_by_tc_id[tc_id] = conf + + tool_trace = [ + entry for entry in (output.get("trace") or []) + if isinstance(entry, dict) and entry.get("node") == "tools" + ] + tool_messages = list(output.get("messages") or []) + + for idx, entry in enumerate(tool_trace): + pending = pending_tool_calls.pop(0) if pending_tool_calls else {} + tool_call_id = pending.get("id") + tool_name = entry.get("tool_name") or pending.get("name") or "unknown" + tool_args = ( + entry.get("tool_args") + if isinstance(entry.get("tool_args"), dict) + else pending.get("args") or {} + ) + + # Check if this tool_call is a confirmation + if tool_call_id and tool_call_id in conf_by_tc_id: + conf = conf_by_tc_id[tool_call_id] + confirmation = conf.get("confirmation", {}) + emitted_confirmation = True + usage_summary["tool_calls"] += 1 + usage_summary["write_tool_calls"] += 1 + usage_summary["evidence_items"] += 1 + tool_latency_ms = int(entry.get("latency_ms", 0) or 0) + usage_summary["tool_latency_ms"] += tool_latency_ms + yield { + "event": "tool_call_confirmation_required", + "data": { + "turnId": request.turn_id, + "agentId": "single-agent", + "agentRole": "single_agent", + "toolCallId": tool_call_id, + "confirmId": confirmation.get("confirm_id", ""), + "action": confirmation.get("tool_name", tool_name), + "description": confirmation.get("description", ""), + "riskLevel": confirmation.get("risk_level", ""), + "permissionExplanation": confirmation.get( + "permission_explanation", "" + ), + "riskExplanation": confirmation.get("risk_explanation", ""), + "affectedResources": confirmation.get( + "affected_resources", [] + ), + "interaction": confirmation.get("interaction", "approval"), + "form": confirmation.get("form"), + "status": "awaiting_confirmation", + "latencyMs": tool_latency_ms, + }, + } + else: + # Normal tool completion + result_status = str(entry.get("result_status") or "unknown").strip().lower() + is_error = result_status == "error" + raw_output = "" + if idx < len(tool_messages): + raw_output = str(getattr(tool_messages[idx], "content", "") or "") + tool_result_summaries.append(f"{tool_name}: {raw_output[:300]}") + usage_summary["tool_calls"] += 1 + usage_summary["read_tool_calls"] += 1 + usage_summary["evidence_items"] += 1 + usage_summary["tool_latency_ms"] += int(entry.get("latency_ms", 0) or 0) + yield { + "event": "tool_call_completed", + "data": { + "turnId": request.turn_id, + "agentId": "single-agent", + "agentRole": "single_agent", + "toolCallId": tool_call_id, + "toolName": tool_name, + "toolArgs": tool_args, + "result": raw_output, + "resultSummary": raw_output[:500], + "status": "error" if is_error else "done", + "isError": is_error, + "latencyMs": entry.get("latency_ms", 0), + }, + } + continue + + # Cancel any orphaned tool_call_started events that never got executed + # (e.g., LLM requested tools but limit was hit before tools_node ran) + for tc in pending_tool_calls: + yield { + "event": "tool_call_completed", + "data": { + "turnId": request.turn_id, + "agentId": "single-agent", + "agentRole": "single_agent", + "toolCallId": tc.get("id"), + "toolName": tc.get("name", "unknown"), + "toolArgs": tc.get("args", {}), + "result": "", + "resultSummary": "已超过单轮工具调用上限，本次调用已取消", + "status": "cancelled", + "isError": False, + }, + } + pending_tool_calls.clear() + + if pending_final_content and not emitted_confirmation: + emitted_final_answer = True + yield { + "event": "final_answer", + "data": { + "turnId": request.turn_id, + "sessionId": request.session_id, + "agentId": "single-agent", + "agentRole": "single_agent", + "content": pending_final_content, + "usageSummary": dict(usage_summary), + }, + } + + if not emitted_final_answer and not emitted_confirmation: + if tool_result_summaries: + details = "\n".join(f"- {item}" for item in tool_result_summaries) + fallback_content = ( + "我已经完成了本轮工具调用，但模型没有正常产出最终答复。" + "已拿到的结果如下：\n" + f"{details}\n\n" + "你可以基于这些结果继续追问，或者重试一次，我会继续推进。" + ) + else: + fallback_content = ( + "本轮执行已结束，但模型没有正常产出最终答复。" + "请直接重试一次；如果问题持续出现，我可以继续帮你定位具体是哪一步卡住了。" + ) + yield { + "event": "final_answer", + "data": { + "turnId": request.turn_id, + "sessionId": request.session_id, + "agentId": "single-agent", + "agentRole": "single_agent", + "content": fallback_content, + "usageSummary": dict(usage_summary), + }, + } + + yield {"event": "done", "data": {"usageSummary": dict(usage_summary)}} diff --git a/crater-agent/crater_agent/tools/__init__.py b/crater-agent/crater_agent/tools/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/crater-agent/crater_agent/tools/definitions.py b/crater-agent/crater_agent/tools/definitions.py new file mode 100644 index 000000000..7ba191fcc --- /dev/null +++ b/crater-agent/crater_agent/tools/definitions.py @@ -0,0 +1,245 @@ +"""Minimal tool definitions for Crater Agent chat.""" + +from typing import Optional + +from langchain_core.tools import tool + + +@tool +def get_job_detail(job_name: str) -> dict: + """获取指定作业的详细状态、资源、时间线和终止信息。""" + pass + + +@tool +def get_job_events(job_name: str) -> dict: + """获取作业关联事件。""" + pass + + +@tool +def get_job_logs(job_name: str, tail: int = 200, keyword: Optional[str] = None) -> dict: + """获取作业日志尾部，可选关键词过滤。""" + pass + + +@tool +def diagnose_job(job_name: str) -> dict: + """对作业进行规则诊断，返回故障分类、证据和建议。""" + pass + + +@tool +def get_diagnostic_context(job_name: str, include_log: bool = True, tail: int = 120) -> dict: + """获取作业诊断上下文，包括元数据、事件、终止状态和可选日志。""" + pass + + +@tool +def search_similar_failures(job_name: str, days: int = 30, limit: int = 5) -> dict: + """搜索与指定作业相似的历史失败案例。""" + pass + + +@tool +def query_job_metrics(job_name: str, metrics: Optional[list[str]] = None, time_range: str = "last_2h") -> dict: + """查询作业 GPU/CPU/内存等监控指标。""" + pass + + +@tool +def analyze_queue_status(job_name: str) -> dict: + """分析 Pending 作业排队或调度原因。""" + pass + + +@tool +def get_realtime_capacity() -> dict: + """读取实时资源容量概览。""" + pass + + +@tool +def list_available_images(limit: int = 30) -> dict: + """列出当前用户可见镜像。""" + pass + + +@tool +def list_available_gpu_models() -> dict: + """列出可选 GPU 型号。""" + pass + + +@tool +def check_quota() -> dict: + """查看当前账户配额摘要。""" + pass + + +@tool +def list_user_jobs(limit: int = 20) -> dict: + """列出当前用户近期作业。""" + pass + + +@tool +def get_job_templates() -> dict: + """列出可创建的作业模板。""" + pass + + +@tool +def get_resource_recommendation(task_type: Optional[str] = None) -> dict: + """根据任务描述推荐基础资源配置。""" + pass + + +@tool +def resubmit_job(job_name: str) -> dict: + """重新提交已有作业；不要先文字询问确认，调用本工具后系统会弹确认卡片。""" + pass + + +@tool +def stop_job(job_name: str) -> dict: + """停止作业；不要先文字询问确认，调用本工具后系统会弹确认卡片。""" + pass + + +@tool +def delete_job(job_name: str) -> dict: + """删除作业；不要先文字询问确认，调用本工具后系统会弹确认卡片。""" + pass + + +@tool +def create_jupyter_job( + name: str, + image_link: str, + cpu: Optional[str] = None, + memory: Optional[str] = None, + gpu_count: Optional[int] = None, + gpu_model: Optional[str] = None, +) -> dict: + """创建 Jupyter 作业；不要先文字询问确认，调用本工具后系统会弹确认卡片。""" + pass + + +@tool +def create_webide_job( + name: str, + image_link: str, + cpu: Optional[str] = None, + memory: Optional[str] = None, + gpu_count: Optional[int] = None, + gpu_model: Optional[str] = None, +) -> dict: + """创建 WebIDE 作业；不要先文字询问确认，调用本工具后系统会弹确认卡片。""" + pass + + +@tool +def create_custom_job( + name: str, + image_link: str, + command: str, + cpu: Optional[str] = None, + memory: Optional[str] = None, + gpu_count: Optional[int] = None, + gpu_model: Optional[str] = None, +) -> dict: + """创建自定义作业；不要先文字询问确认，调用本工具后系统会弹确认卡片。""" + pass + + +@tool +def create_pytorch_job( + name: str, + image_link: str, + command: Optional[str] = None, + cpu: Optional[str] = None, + memory: Optional[str] = None, + gpu_count: Optional[int] = None, + gpu_model: Optional[str] = None, +) -> dict: + """创建 PyTorch 作业；不要先文字询问确认，调用本工具后系统会弹确认卡片。""" + pass + + +@tool +def create_tensorflow_job( + name: str, + image_link: str, + command: Optional[str] = None, + cpu: Optional[str] = None, + memory: Optional[str] = None, + gpu_count: Optional[int] = None, + gpu_model: Optional[str] = None, +) -> dict: + """创建 TensorFlow 作业；不要先文字询问确认，调用本工具后系统会弹确认卡片。""" + pass + + +AUTO_TOOLS = [ + get_job_detail, + get_job_events, + get_job_logs, + diagnose_job, + get_diagnostic_context, + search_similar_failures, + query_job_metrics, + analyze_queue_status, + get_realtime_capacity, + list_available_images, + list_available_gpu_models, + check_quota, + list_user_jobs, + get_job_templates, + get_resource_recommendation, +] + +AUTO_ACTION_TOOLS = [] + +CONFIRM_TOOLS = [ + resubmit_job, + stop_job, + delete_job, + create_jupyter_job, + create_webide_job, + create_custom_job, + create_pytorch_job, + create_tensorflow_job, +] + +ALL_TOOLS = AUTO_TOOLS + AUTO_ACTION_TOOLS + CONFIRM_TOOLS + +DEPRECATED_TOOL_NAMES: set[str] = set() +INTERNAL_TOOLS = [] +INTERNAL_TOOL_NAMES = set() +READ_ONLY_TOOL_NAMES = {tool_item.name for tool_item in AUTO_TOOLS} +AUTO_ACTION_TOOL_NAMES = set() +CONFIRM_TOOL_NAMES = {tool_item.name for tool_item in CONFIRM_TOOLS} +WRITE_TOOL_NAMES = CONFIRM_TOOL_NAMES +ADMIN_ONLY_TOOL_NAMES = set() + +ROLE_ALLOWED_TOOL_NAMES = { + "planner": READ_ONLY_TOOL_NAMES, + "coordinator": READ_ONLY_TOOL_NAMES, + "explorer": READ_ONLY_TOOL_NAMES, + "executor": {tool_item.name for tool_item in ALL_TOOLS}, + "verifier": READ_ONLY_TOOL_NAMES, + "guide": set(), + "general": READ_ONLY_TOOL_NAMES, + "single_agent": {tool_item.name for tool_item in ALL_TOOLS}, +} + + +def is_tool_allowed_for_role(role: Optional[str], tool_name: str) -> bool: + normalized_role = (role or "single_agent").strip().lower() or "single_agent" + allowed = ROLE_ALLOWED_TOOL_NAMES.get(normalized_role, ROLE_ALLOWED_TOOL_NAMES["single_agent"]) + return tool_name in allowed + + +def is_actor_allowed_for_tool(actor_role: Optional[str], tool_name: str) -> bool: + return True diff --git a/crater-agent/crater_agent/tools/executor.py b/crater-agent/crater_agent/tools/executor.py new file mode 100644 index 000000000..bf2670b58 --- /dev/null +++ b/crater-agent/crater_agent/tools/executor.py @@ -0,0 +1,218 @@ +"""Tool executor that calls the Crater Go backend. + +All tool executions are proxied through Go's /v1/agent/tools/execute endpoint, +which handles permission checks, data access, and audit logging. +""" + +from __future__ import annotations + +import json +import logging +import time +from typing import Any, Protocol + +logger = logging.getLogger(__name__) + +import httpx + +from crater_agent.config import settings +from crater_agent.tools.definitions import ( + is_actor_allowed_for_tool, + is_tool_allowed_for_role, +) + + +def normalize_tool_args_for_backend(tool_name: str, tool_args: dict[str, Any]) -> dict[str, Any]: + """Normalize provider-specific argument variants before backend dispatch.""" + normalized = dict(tool_args or {}) + if ( + tool_name == "notify_job_owner" + and "job_names" not in normalized + and normalized.get("job_name") + ): + normalized["job_names"] = normalized["job_name"] + return normalized + + +class ToolExecutorProtocol(Protocol): + """Protocol for production tool executors.""" + + async def execute( + self, + tool_name: str, + tool_args: dict[str, Any], + session_id: str, + user_id: int, + turn_id: str | None = None, + tool_call_id: str | None = None, + agent_id: str | None = None, + agent_role: str | None = None, + actor_role: str | None = None, + execution_backend: str | None = None, + ) -> dict[str, Any]: ... + + +class GoBackendToolExecutor: + """Executes tools by calling the Crater Go backend.""" + + def __init__(self, backend_url: str | None = None): + self.backend_url = backend_url or settings.crater_backend_url + self.client = httpx.AsyncClient( + base_url=self.backend_url, + timeout=settings.tool_execution_timeout, + ) + + async def execute( + self, + tool_name: str, + tool_args: dict[str, Any], + session_id: str, + user_id: int, + turn_id: str | None = None, + tool_call_id: str | None = None, + agent_id: str | None = None, + agent_role: str | None = None, + actor_role: str | None = None, + execution_backend: str | None = None, + ) -> dict[str, Any]: + """Execute a tool via Go backend HTTP call. + + Returns: + For auto tools: {"status": "success", "result": {...}} + For confirm tools: {"status": "confirmation_required", "confirm_id": "...", ...} + On error: {"status": "error", "message": "..."} + """ + start_time = time.monotonic() + tool_args = normalize_tool_args_for_backend(tool_name, tool_args) + logger.info("[tool] execute: %s args=%s session=%s actor=%s", tool_name, tool_args, session_id, actor_role) + normalized_role = (agent_role or "single_agent").strip().lower() or "single_agent" + if not is_tool_allowed_for_role(normalized_role, tool_name): + return { + "status": "error", + "error_type": "tool_policy", + "retryable": False, + "message": f"Tool {tool_name} is not allowed for agent role {normalized_role}", + "_latency_ms": int((time.monotonic() - start_time) * 1000), + } + if not is_actor_allowed_for_tool(actor_role, tool_name): + return { + "status": "error", + "error_type": "tool_policy", + "retryable": False, + "message": f"你当前没有管理员权限，不能执行 {tool_name}；如确需处理，请联系平台管理员或切换到管理员页面后再操作。", + "_latency_ms": int((time.monotonic() - start_time) * 1000), + } + try: + request_body: dict[str, Any] = { + "tool_name": tool_name, + "tool_args": tool_args, + "session_id": session_id, + "turn_id": turn_id, + "tool_call_id": tool_call_id, + "agent_id": agent_id, + "agent_role": normalized_role, + } + if execution_backend: + request_body["execution_backend"] = execution_backend + # System-level agents (e.g. approval evaluator) don't have a real + # AgentSession in the database. Pass internal_context so the Go + # backend resolves an admin token directly instead of doing a + # session lookup that would fail with "session not found". + if actor_role == "system": + request_body["internal_context"] = { + "role": "admin", + "username": "agent-approval", + } + + resp = await self.client.post( + "/api/agent/tools/execute", + headers={ + "X-Agent-Internal-Token": settings.crater_backend_internal_token, + }, + json=request_body, + ) + resp.raise_for_status() + payload = resp.json() + result = payload.get("data", payload) + if not isinstance(result, dict): + result = {"status": "success", "result": result} + result["_latency_ms"] = int((time.monotonic() - start_time) * 1000) + status = result.get("status", "ok") + if status == "error": + logger.warning("[tool] %s failed: %s (latency=%dms)", tool_name, result.get("message") or result.get("error", ""), result["_latency_ms"]) + else: + # Truncate result summary to avoid flooding logs + brief = str(result.get("result", ""))[:200] + logger.info("[tool] %s ok: %s (latency=%dms)", tool_name, brief, result["_latency_ms"]) + return result + except httpx.TimeoutException: + logger.warning("[tool] %s timed out after %ds", tool_name, settings.tool_execution_timeout) + return { + "status": "error", + "error_type": "timeout", + "retryable": True, + "message": f"Tool {tool_name} 执行超时 ({settings.tool_execution_timeout}s)", + "_latency_ms": int((time.monotonic() - start_time) * 1000), + } + except httpx.HTTPStatusError as e: + detail = "" + try: + body = e.response.json() + if isinstance(body, dict): + detail = str(body.get("msg") or body.get("message") or body.get("error") or "") + except (json.JSONDecodeError, ValueError, TypeError): + detail = "" + + status_code = e.response.status_code + if status_code in (401, 403): + error_type = "auth" + retryable = False + elif status_code == 404: + error_type = "not_found" + retryable = False + elif status_code == 429: + error_type = "rate_limit" + retryable = True + elif 500 <= status_code < 600: + error_type = "server" + retryable = True + else: + error_type = "http" + retryable = False + + logger.warning("[tool] %s HTTP error: %d %s", tool_name, status_code, detail) + return { + "status": "error", + "error_type": error_type, + "retryable": retryable, + "status_code": status_code, + "message": ( + f"Tool {tool_name} 执行失败: HTTP {status_code}" + + (f" - {detail}" if detail else "") + ), + "_latency_ms": int((time.monotonic() - start_time) * 1000), + } + except httpx.RequestError as e: + logger.warning("[tool] %s network error: %s", tool_name, e) + return { + "status": "error", + "error_type": "network", + "retryable": True, + "message": f"Tool {tool_name} 网络异常: {e}", + "_latency_ms": int((time.monotonic() - start_time) * 1000), + } + except Exception as e: + return { + "status": "error", + "error_type": "unexpected", + "retryable": False, + "message": f"Tool {tool_name} 执行异常: {e}", + "_latency_ms": int((time.monotonic() - start_time) * 1000), + } + + async def close(self): + await self.client.aclose() + + +class CompositeToolExecutor(GoBackendToolExecutor): + """Compatibility alias: all tool execution now goes through Go backend.""" diff --git a/crater-agent/crater_agent/tools/tool_selector.py b/crater-agent/crater_agent/tools/tool_selector.py new file mode 100644 index 000000000..4789fae58 --- /dev/null +++ b/crater-agent/crater_agent/tools/tool_selector.py @@ -0,0 +1,140 @@ +"""Role-based tool selector and capability sanitization. + +Security model: +- Privilege elevation must come from trusted actor identity, never from page + route, URL, or user-provided text. +- Page scope may reduce visibility (for example an admin browsing a user page), + but it must never expand a normal user into admin scope. +- Capability payloads are sanitized before being exposed to prompts so an + injected or malformed tool catalog cannot leak admin-only tools. +""" + +from __future__ import annotations + +import logging +from typing import Any + +from crater_agent.tools.definitions import ADMIN_ONLY_TOOL_NAMES, is_actor_allowed_for_tool + +logger = logging.getLogger(__name__) + +_ADMIN_ROLES = frozenset({"roleadmin", "admin", "platform_admin", "system_admin"}) +_SYSTEM_ROLES = frozenset({"system"}) + + +def normalize_actor_role(actor_role: Any) -> str: + return str(actor_role or "user").strip().lower() or "user" + + +def is_privileged_actor_role(actor_role: Any) -> bool: + normalized = normalize_actor_role(actor_role) + return normalized in _ADMIN_ROLES or normalized in _SYSTEM_ROLES + + +def _infer_page_scope(context: dict[str, Any]) -> str: + capabilities = context.get("capabilities") or {} + surface = capabilities.get("surface") if isinstance(capabilities, dict) else {} + if isinstance(surface, dict): + declared_scope = str(surface.get("page_scope") or "").strip().lower() + if declared_scope in {"admin", "user"}: + return declared_scope + + page = context.get("page") or {} + route = str(page.get("route") or "").strip().lower() + url = str(page.get("url") or "").strip().lower() + if route.startswith("/admin") or "/admin/" in route or url.startswith("/admin") or "/admin/" in url: + return "admin" + if route or url: + return "user" + return "" + + +def _resolve_actor_role(context: dict[str, Any]) -> str: + """Determine the effective actor role from trusted identity plus page scope. + + Trusted actor identity (JWT/backend context) is the only source that can + grant admin privileges. Page scope can narrow visibility, but never widen + it for a normal user. + """ + actor = context.get("actor") or {} + trusted_role = normalize_actor_role(actor.get("role")) + if trusted_role in _SYSTEM_ROLES: + return trusted_role + + page_scope = _infer_page_scope(context) + if page_scope == "admin" and not is_privileged_actor_role(trusted_role): + return "user" + return trusted_role + + +def sanitize_enabled_tool_names(context: dict[str, Any], enabled_tool_names: list[Any] | None) -> list[str]: + effective_role = _resolve_actor_role(context) + sanitized: list[str] = [] + seen: set[str] = set() + for item in enabled_tool_names or []: + tool_name = str(item or "").strip() + if not tool_name or tool_name in seen: + continue + if not is_actor_allowed_for_tool(effective_role, tool_name): + continue + sanitized.append(tool_name) + seen.add(tool_name) + return sanitized + + +def sanitize_capabilities_for_context( + context: dict[str, Any], + capabilities: dict[str, Any] | None, +) -> dict[str, Any]: + raw = dict(capabilities or {}) + sanitized = dict(raw) + + effective_role = _resolve_actor_role(context) + sanitized_enabled = sanitize_enabled_tool_names(context, raw.get("enabled_tools")) + sanitized["enabled_tools"] = sanitized_enabled + + raw_catalog = raw.get("tool_catalog") + filtered_catalog: list[dict[str, Any]] = [] + enabled_set = set(sanitized_enabled) + if isinstance(raw_catalog, list): + for item in raw_catalog: + if not isinstance(item, dict): + continue + name = str(item.get("name") or "").strip() + if not name: + continue + if enabled_set and name not in enabled_set: + continue + if not is_actor_allowed_for_tool(effective_role, name): + continue + filtered_catalog.append(dict(item)) + sanitized["tool_catalog"] = filtered_catalog + + surface = dict(raw.get("surface") or {}) + page_scope = _infer_page_scope({"capabilities": {"surface": surface}, "page": context.get("page")}) + surface["page_scope"] = "admin" if is_privileged_actor_role(effective_role) else "user" + sanitized["surface"] = surface + return sanitized + + +def select_tools_for_context(context: dict[str, Any], all_tools: list) -> list: + """Select tools based on actor role. + + - Admin: returns all tools (no filtering) + - User: returns only user-level tools + + This replaces the previous URL/route-based filtering. + """ + role = _resolve_actor_role(context) + + if role in _ADMIN_ROLES or role in _SYSTEM_ROLES: + logger.debug("Tool selector: admin role %r → all %d tools", role, len(all_tools)) + return all_tools + + # Regular user: expose every tool that is not admin-only. + filtered = [t for t in all_tools if t.name not in ADMIN_ONLY_TOOL_NAMES] + logger.debug( + "Tool selector: user role %r → %d/%d tools", + role, len(filtered), len(all_tools), + ) + return filtered diff --git a/crater-agent/pyproject.toml b/crater-agent/pyproject.toml new file mode 100644 index 000000000..3f498d9c1 --- /dev/null +++ b/crater-agent/pyproject.toml @@ -0,0 +1,44 @@ +[project] +name = "crater-agent" +version = "0.1.0" +description = "Intelligent Agent for Crater AI Platform - ReAct-based autonomous diagnosis and operations assistant" +requires-python = ">=3.11" +dependencies = [ + "langgraph>=0.2.0", + "langchain-core>=0.3.0", + "langchain-openai>=0.2.0", + "fastapi>=0.115.0", + "uvicorn[standard]>=0.32.0", + "httpx>=0.27.0", + "pyyaml>=6.0", + "pydantic>=2.0", + "pydantic-settings>=2.0", + "sse-starlette>=2.0", + "camel-ai>=0.2.16", + "ddgs>=0.1", +] + +[build-system] +requires = ["setuptools>=68.0"] +build-backend = "setuptools.build_meta" + +[tool.setuptools.packages.find] +include = ["crater_agent*"] + +[project.optional-dependencies] +dev = [ + "pytest>=8.0", + "pytest-asyncio>=0.24.0", + "ruff>=0.8.0", +] + +[tool.ruff] +line-length = 100 +target-version = "py311" + +[tool.ruff.lint] +select = ["E", "F", "I", "W"] + +[tool.pytest.ini_options] +asyncio_mode = "auto" +testpaths = ["tests"] \ No newline at end of file diff --git a/frontend/docs/images/model-metadata-governance-detail.jpg b/frontend/docs/images/model-metadata-governance-detail.jpg deleted file mode 100644 index 3de6f266b..000000000 Binary files a/frontend/docs/images/model-metadata-governance-detail.jpg and /dev/null differ diff --git a/frontend/docs/images/model-metadata-governance-list.jpg b/frontend/docs/images/model-metadata-governance-list.jpg deleted file mode 100644 index da3ada7a1..000000000 Binary files a/frontend/docs/images/model-metadata-governance-list.jpg and /dev/null differ diff --git a/frontend/pnpm-workspace.yaml b/frontend/pnpm-workspace.yaml new file mode 100644 index 000000000..d3199a652 --- /dev/null +++ b/frontend/pnpm-workspace.yaml @@ -0,0 +1,11 @@ +allowBuilds: + '@swc/core': true + '@tailwindcss/oxide': true + esbuild: true + msw: true + +onlyBuiltDependencies: + - '@swc/core' + - '@tailwindcss/oxide' + - esbuild + - msw diff --git a/frontend/src/components/aiops/AIChatAssistantProvider.tsx b/frontend/src/components/aiops/AIChatAssistantProvider.tsx new file mode 100644 index 000000000..83eda1eb4 --- /dev/null +++ b/frontend/src/components/aiops/AIChatAssistantProvider.tsx @@ -0,0 +1,32 @@ +'use client' + +import { useState } from 'react' + +import { AIChatDrawer } from '@/components/aiops/AIChatDrawer' +import { FloatingAssistantButton } from '@/components/aiops/FloatingAssistantButton' + +export function AIChatAssistantProvider({ + children, + currentJobName, +}: { + children: React.ReactNode + currentJobName?: string +}) { + const [isChatOpen, setIsChatOpen] = useState(false) + + return ( + <> + {children} + + {/* Floating Button - Always visible */} + setIsChatOpen(true)} /> + + {/* Chat Drawer - Opens on click */} + setIsChatOpen(false)} + currentJobName={currentJobName} + /> + + ) +} diff --git a/frontend/src/components/aiops/AIChatDrawer.tsx b/frontend/src/components/aiops/AIChatDrawer.tsx new file mode 100644 index 000000000..316baa9f7 --- /dev/null +++ b/frontend/src/components/aiops/AIChatDrawer.tsx @@ -0,0 +1,3199 @@ +'use client' + +import { useQuery, useQueryClient } from '@tanstack/react-query' +import { + AlertCircle, + CheckCircle, + Copy, + HelpCircle, + History, + Loader2, + PanelLeftClose, + Pencil, + Pin, + Plus, + RotateCcw, + Send, + Sparkles, + Square, + Trash2, + Users, + X, + Zap, +} from 'lucide-react' +import { useCallback, useEffect, useRef, useState } from 'react' +import { useTranslation } from 'react-i18next' +import ReactMarkdown from 'react-markdown' +import remarkGfm from 'remark-gfm' +import { toast } from 'sonner' + +import { + AlertDialog, + AlertDialogAction, + AlertDialogCancel, + AlertDialogContent, + AlertDialogDescription, + AlertDialogFooter, + AlertDialogHeader, + AlertDialogTitle, +} from '@/components/ui/alert-dialog' +import { Badge } from '@/components/ui/badge' +import { Button } from '@/components/ui/button' +import { Input } from '@/components/ui/input' +import { ScrollArea } from '@/components/ui/scroll-area' +import { + Select, + SelectContent, + SelectItem, + SelectTrigger, + SelectValue, +} from '@/components/ui/select' +import { Textarea } from '@/components/ui/textarea' + +import { + apiDeleteSession, + apiGetAgentConfigSummary, + apiGetSessionMessages, + apiGetSessionToolCalls, + apiGetSessionTurns, + apiGetTurnEvents, + apiListFeedbacks, + apiListSessions, + apiParameterUpdate, + apiPinSession, + apiRenameSession, + connectAgentAskStream, + connectAgentChat, + connectAgentResume, +} from '@/services/api/agent' +import type { + AgentConfigSummary, + AgentConfirmationForm, + AgentEvent, + AgentFeedback, + AgentMessage, + AgentSession, + AgentSurface, + AgentToolCall, + AgentTurn, + ParameterReviewPayload, +} from '@/services/api/agent' +import type { AgentSSEEvent } from '@/services/api/agent' + +import { cn } from '@/lib/utils' + +import { ConfirmActionCard } from './ConfirmActionCard' +import { FeedbackCard } from './FeedbackCard' +import { ParameterReviewCard } from './ParameterReviewCard' +import { ThinkingIndicator } from './ThinkingIndicator' +import { ToolCallCard } from './ToolCallCard' + +const AGENT_INPUT_MAX_LENGTH = 4000 +const AGENT_LAST_SESSION_STORAGE_KEY = 'crater-agent-last-session-id' +const CHAT_DRAWER_DEFAULT_WIDTH = 680 +const CHAT_DRAWER_MIN_WIDTH = 440 +const CHAT_DRAWER_MAX_WIDTH = 920 +const CHAT_INPUT_MAX_HEIGHT = 144 + +type AgentEntryPoint = 'default' | 'node_analysis' +type ChatMode = 'llm' | 'agent' + +type ConversationItemKind = + | 'user' + | 'thinking' + | 'message' + | 'tool_call' + | 'confirmation_required' + | 'parameter_review' + | 'error' + +interface ConversationItem { + id: string + kind: ConversationItemKind + /** For 'user' and 'message' kinds */ + text?: string + requestId?: string + requestState?: 'running' | 'done' | 'awaiting_confirmation' | 'failed' | 'cancelled' + requestError?: string + requestSessionId?: string | null + requestOrchestrationMode?: 'single_agent' | 'ask' + /** For 'thinking' — may be partial/streaming */ + thinkingContent?: string + /** For 'tool_call' (single_agent mode only) */ + toolName?: string + toolArgs?: Record + toolStatus?: 'executing' | 'awaiting_confirmation' | 'done' | 'error' | 'cancelled' + toolResult?: string + /** For 'confirmation_required' */ + confirmId?: string + confirmToolCallId?: string + confirmAction?: string + confirmDescription?: string + confirmRiskLevel?: string + confirmPermissionExplanation?: string + confirmRiskExplanation?: string + confirmAffectedResources?: string[] + confirmInteraction?: string + confirmForm?: AgentConfirmationForm + retryRequestId?: string + /** For agent_event in single_agent (thinking update) */ + agentRole?: string + /** For 'parameter_review' */ + parameterReview?: ParameterReviewPayload + /** For feedback — the DB message id or turn id to use as targetId */ + feedbackTargetId?: string + timestamp: Date +} + +interface AgentPendingRequest { + requestId: string + sessionId: string | null + message: string + orchestrationMode: 'single_agent' + pageContext: { + route?: string + url: string + jobName?: string + jobStatus?: string + nodeName?: string + entryPoint?: AgentEntryPoint + surface?: AgentSurface + } + clientContext?: { + locale?: string + timezone?: string + } +} + +interface ActiveAgentRequestState { + requestId: string + hasFinalResponse: boolean + awaitingConfirmation: boolean +} + +interface InterruptConfirmState { + title: string + description: string + confirmLabel: string +} + +interface SessionDeleteConfirmState { + sessionId: string + title: string +} + +interface AgentEventPayload { + turnId?: string + agentId?: string + agentRole?: string + parentAgentId?: string + targetAgentId?: string + targetAgentRole?: string + summary?: string + status?: string + title?: string + eventType?: string + verificationResult?: string + content?: string + toolName?: string + name?: string + tool?: string + toolArgs?: Record + args?: Record + arguments?: Record + toolCallId?: string + id?: string + resultSummary?: string + result?: unknown + isError?: boolean + partial?: boolean + sessionId?: string + feedbackTargetId?: string + confirmId?: string + confirm_id?: string + action?: string + tool_name?: string + description?: string + riskLevel?: string + risk_level?: string + permissionExplanation?: string + permission_explanation?: string + riskExplanation?: string + risk_explanation?: string + affectedResources?: string[] + affected_resources?: string[] + interaction?: string + form?: AgentConfirmationForm +} + +// ───────────────────────────────────────────────────────────────────────────── + +interface AIChatDrawerProps { + isOpen: boolean + onClose: () => void + currentJobName?: string +} + +function parseAgentJSON(value: unknown): Record | null { + if (!value) return null + if (typeof value === 'object') return value as Record + if (typeof value !== 'string') return null + try { + return JSON.parse(value) as Record + } catch { + return null + } +} + +function generateAgentRequestId() { + if (typeof crypto !== 'undefined' && typeof crypto.randomUUID === 'function') { + return crypto.randomUUID() + } + return `agent-req-${Date.now()}-${Math.random().toString(36).slice(2, 10)}` +} + +function resizeChatTextarea(textarea: HTMLTextAreaElement | null) { + if (!textarea) return + textarea.style.height = 'auto' + const nextHeight = Math.min(textarea.scrollHeight, CHAT_INPUT_MAX_HEIGHT) + textarea.style.height = `${nextHeight}px` + textarea.style.overflowY = textarea.scrollHeight > CHAT_INPUT_MAX_HEIGHT ? 'auto' : 'hidden' +} + +function normalizeAgentOrchestrationMode( + _mode: 'single_agent' | 'ask' | undefined +): 'single_agent' { + void _mode + return 'single_agent' +} + +function inferAgentEntryPoint(pathname: string): AgentEntryPoint { + if (/\/admin\/nodes(\/|$)/.test(pathname)) { + return 'node_analysis' + } + return 'default' +} + +function getToolStatusFromResult(resultStatus: string): ConversationItem['toolStatus'] { + switch (resultStatus) { + case 'success': + return 'done' + case 'cancelled': + case 'rejected': + return 'cancelled' + case 'await_confirm': + case 'confirmation_required': + return 'awaiting_confirmation' + default: + return 'error' + } +} + +function getMessageRequestId(message: AgentMessage): string | undefined { + const metadata = parseAgentJSON(message.metadata) + const requestId = metadata?.requestId + return typeof requestId === 'string' ? requestId : undefined +} + +function getRequestStateFromTurnStatus( + status: string | undefined +): ConversationItem['requestState'] | undefined { + switch (status) { + case 'running': + return 'running' + case 'completed': + return 'done' + case 'awaiting_confirmation': + return 'awaiting_confirmation' + case 'failed': + return 'failed' + default: + return undefined + } +} + +function getFailedRequestMessage() { + return 'Agent 未返回最终答复，执行可能已中断。' +} + +function getTurnFailureMessage(turn: AgentTurn | undefined): string | undefined { + if (!turn) return undefined + const metadata = parseAgentJSON(turn.metadata) + const errorMessage = metadata?.errorMessage + return typeof errorMessage === 'string' && errorMessage.trim() ? errorMessage : undefined +} + +function getConversationItemSortWeight(item: ConversationItem): number { + switch (item.kind) { + case 'user': + return 0 + case 'thinking': + return 1 + case 'tool_call': + return 2 + case 'confirmation_required': + return 3 + case 'message': + return 4 + case 'error': + return 5 + default: + return 99 + } +} + +function formatAgentSessionDate(dateString: string): string { + const date = new Date(dateString) + if (Number.isNaN(date.getTime())) { + return dateString + } + return `${date.getFullYear()}年${date.getMonth() + 1}月${date.getDate()}日` +} + +function mapSingleAgentRunEventsToToolItems( + turns: AgentTurn[], + toolCalls: AgentToolCall[], + runEventsByTurn: Record +): ConversationItem[] { + const persistedToolCallIds = new Set( + toolCalls.map((toolCall) => String(toolCall.toolCallId || '').trim()).filter(Boolean) + ) + + const itemsById = new Map() + + for (const turn of turns) { + const turnEvents = runEventsByTurn[turn.turnId] ?? [] + for (const event of turnEvents) { + if ( + event.eventType !== 'tool_call_started' && + event.eventType !== 'tool_call_completed' && + event.eventType !== 'tool_call_confirmation_required' + ) { + continue + } + + const metadata = parseAgentJSON(event.metadata) + const toolCallId = String(metadata?.toolCallId || '').trim() + if (!toolCallId || persistedToolCallIds.has(toolCallId)) continue + + const itemId = `toolcall-${toolCallId}` + const previous = itemsById.get(itemId) + const toolName = + typeof metadata?.toolName === 'string' && metadata.toolName.trim() + ? metadata.toolName + : event.title || 'unknown' + const toolArgs = + (typeof metadata?.toolArgs === 'object' && metadata.toolArgs !== null + ? (metadata.toolArgs as Record) + : typeof metadata?.args === 'object' && metadata.args !== null + ? (metadata.args as Record) + : typeof metadata?.arguments === 'object' && metadata.arguments !== null + ? (metadata.arguments as Record) + : previous?.toolArgs) ?? {} + + let toolStatus: ConversationItem['toolStatus'] = previous?.toolStatus ?? 'executing' + if (event.eventType === 'tool_call_started') { + toolStatus = 'executing' + } else if (event.eventType === 'tool_call_confirmation_required') { + toolStatus = 'awaiting_confirmation' + } else if (event.eventStatus === 'error' || metadata?.isError === true) { + toolStatus = 'error' + } else { + toolStatus = 'done' + } + + const resultValue = + typeof metadata?.resultSummary === 'string' + ? metadata.resultSummary + : typeof metadata?.result === 'string' + ? metadata.result + : event.content || previous?.toolResult + + itemsById.set(itemId, { + id: itemId, + kind: 'tool_call', + toolName, + toolArgs, + toolStatus, + toolResult: resultValue, + timestamp: previous?.timestamp ?? new Date(event.startedAt || event.createdAt), + }) + } + } + + return Array.from(itemsById.values()) +} + +function mapSessionHistoryToConversationItems( + messages: AgentMessage[], + toolCalls: AgentToolCall[], + turns: AgentTurn[], + runEventsByTurn: Record +): ConversationItem[] { + const latestTurnByRequestId = new Map() + for (const turn of turns) { + if (!turn.requestId) continue + const existing = latestTurnByRequestId.get(turn.requestId) + if (!existing || new Date(turn.startedAt).getTime() >= new Date(existing.startedAt).getTime()) { + latestTurnByRequestId.set(turn.requestId, turn) + } + } + + const messageItems = messages.map((message) => { + const requestId = message.role === 'user' ? getMessageRequestId(message) : undefined + const turn = requestId ? latestTurnByRequestId.get(requestId) : undefined + const requestState = + message.role === 'user' ? getRequestStateFromTurnStatus(turn?.status) : undefined + const failureMessage = turn ? getTurnFailureMessage(turn) : undefined + + return { + id: `history-${message.id}`, + kind: message.role === 'user' ? 'user' : 'message', + text: message.content, + requestId, + requestState, + requestError: + requestState === 'failed' ? (failureMessage ?? getFailedRequestMessage()) : undefined, + requestSessionId: message.sessionId ?? turn?.sessionId ?? null, + requestOrchestrationMode: turn?.orchestrationMode, + feedbackTargetId: message.role === 'assistant' ? String(message.id) : undefined, + timestamp: new Date(message.createdAt), + } satisfies ConversationItem + }) + + const toolItems = toolCalls.flatMap((toolCall) => { + const toolArgs = parseAgentJSON(toolCall.toolArgs) ?? {} + const baseTimestamp = new Date(toolCall.createdAt) + const itemId = `history-tool-${toolCall.id}` + const resultSummary = + typeof toolCall.toolResult === 'string' + ? toolCall.toolResult + : JSON.stringify(toolCall.toolResult ?? '') + const items: ConversationItem[] = [ + { + id: itemId, + kind: 'tool_call', + toolName: toolCall.toolName, + toolArgs, + toolStatus: getToolStatusFromResult(toolCall.resultStatus), + toolResult: resultSummary, + timestamp: baseTimestamp, + }, + ] + return items + }) + + const derivedSingleAgentToolItems = mapSingleAgentRunEventsToToolItems( + turns, + toolCalls, + runEventsByTurn + ) + + const confirmationItems = toolCalls.flatMap((toolCall) => { + if ( + toolCall.resultStatus !== 'await_confirm' && + toolCall.resultStatus !== 'confirmation_required' + ) { + return [] + } + + const toolResult = parseAgentJSON(toolCall.toolResult) + + return [ + { + id: `history-confirm-${toolCall.id}`, + kind: 'confirmation_required' as const, + confirmId: String(toolCall.id), + confirmToolCallId: toolCall.turnId + ? toolCall.toolCallId + ? `toolcall-${toolCall.toolCallId}` + : `history-tool-${toolCall.id}` + : `history-tool-${toolCall.id}`, + confirmAction: toolCall.toolName, + confirmDescription: (toolResult?.description as string) ?? `等待确认 ${toolCall.toolName}`, + confirmRiskLevel: (toolResult?.riskLevel as string) ?? (toolResult?.risk_level as string), + confirmPermissionExplanation: + (toolResult?.permissionExplanation as string) ?? + (toolResult?.permission_explanation as string), + confirmRiskExplanation: + (toolResult?.riskExplanation as string) ?? (toolResult?.risk_explanation as string), + confirmAffectedResources: + (toolResult?.affectedResources as string[]) ?? + (toolResult?.affected_resources as string[]), + confirmInteraction: (toolResult?.interaction as string) ?? 'approval', + confirmForm: (toolResult?.form as AgentConfirmationForm) ?? undefined, + timestamp: new Date(new Date(toolCall.createdAt).getTime() + 1), + } satisfies ConversationItem, + ] + }) + + return [...messageItems, ...toolItems, ...derivedSingleAgentToolItems, ...confirmationItems].sort( + (a, b) => { + const timestampDiff = a.timestamp.getTime() - b.timestamp.getTime() + if (timestampDiff !== 0) return timestampDiff + return getConversationItemSortWeight(a) - getConversationItemSortWeight(b) + } + ) +} + +// ── Markdown renderer components ────────────────────────────────────────────── + +const markdownComponents = { + p: ({ children }: { children?: React.ReactNode }) => ( +

{children}

+ ), + strong: ({ children }: { children?: React.ReactNode }) => ( + {children} + ), + ul: ({ children }: { children?: React.ReactNode }) => ( +
    + {children} +
+ ), + ol: ({ children }: { children?: React.ReactNode }) => ( +
    + {children} +
+ ), + li: ({ children }: { children?: React.ReactNode }) => ( +
  • {children}
    • + ), + pre: ({ children }: { children?: React.ReactNode }) => ( + 
    {children}
    + ), + code: ({ children, className }: { children?: React.ReactNode; className?: string }) => { + const isInline = !className + return isInline ? ( + + {children} + + ) : ( + {children} + ) + }, + h1: ({ children }: { children?: React.ReactNode }) => ( +

    {children}

    + ), + h2: ({ children }: { children?: React.ReactNode }) => ( +

    {children}

    + ), + h3: ({ children }: { children?: React.ReactNode }) => ( +

    {children}

    + ), + h4: ({ children }: { children?: React.ReactNode }) => ( +

    {children}

    + ), + blockquote: ({ children }: { children?: React.ReactNode }) => ( +
    + {children} +
    + ), + table: ({ children }: { children?: React.ReactNode }) => ( +
    + {children}
    +
    + ), + th: ({ children }: { children?: React.ReactNode }) => ( + + {children} + + ), + td: ({ children }: { children?: React.ReactNode }) => ( + + {children} + + ), +} + +// ── Main Component ──────────────────────────────────────────────────────────── + +export function AIChatDrawer({ isOpen, onClose, currentJobName }: AIChatDrawerProps) { + const { t } = useTranslation() + const queryClient = useQueryClient() + const [input, setInput] = useState('') + const [showHelp, setShowHelp] = useState(false) + const [chatMode, setChatMode] = useState('llm') + const [drawerWidth, setDrawerWidth] = useState(CHAT_DRAWER_DEFAULT_WIDTH) + const [editingAgentMessageId, setEditingAgentMessageId] = useState(null) + const [editingAgentDraft, setEditingAgentDraft] = useState('') + const isAdminRoute = + typeof window !== 'undefined' && window.location.pathname.startsWith('/admin') + const agentSurface: AgentSurface = isAdminRoute ? 'admin' : 'user' + const agentSessionStorageKey = `${AGENT_LAST_SESSION_STORAGE_KEY}:${agentSurface}` + const messagesEndRef = useRef(null) + const agentInputRef = useRef(null) + + // ── Agent mode state ────────────────────────────────────────────────────── + const [conversationItems, setConversationItems] = useState([]) + const [agentStreaming, setAgentStreaming] = useState(false) + const [pendingConfirmIds, setPendingConfirmIds] = useState([]) + const [agentHistoryLoading, setAgentHistoryLoading] = useState(false) + const [agentHistoryError, setAgentHistoryError] = useState(null) + const [selectedAgentSessionId, setSelectedAgentSessionId] = useState(null) + const [feedbackMap, setFeedbackMap] = useState>({}) + const [orchestrationMode, setOrchestrationMode] = useState<'single_agent'>('single_agent') + const [retryableAgentRequest, setRetryableAgentRequest] = useState( + null + ) + const [failedAgentRequests, setFailedAgentRequests] = useState< + Record + >({}) + const [interruptConfirmState, setInterruptConfirmState] = useState( + null + ) + const [sessionDeleteConfirmState, setSessionDeleteConfirmState] = + useState(null) + const [sessionActionLoading, setSessionActionLoading] = useState<{ + sessionId: string + action: 'pin' | 'delete' | 'rename' + } | null>(null) + const [renameState, setRenameState] = useState<{ + sessionId: string + draftTitle: string + } | null>(null) + const renameInputRef = useRef(null) + const [sessionPanelOpen, setSessionPanelOpen] = useState(true) + const agentAbortRef = useRef(null) + const agentHistoryRequestIdRef = useRef(0) + const agentInteractionVersionRef = useRef(0) + const lastAgentHistorySessionIdRef = useRef(null) + const lastLoadedAgentSessionIdRef = useRef(null) + const lastAgentRequestRef = useRef(null) + const activeAgentRequestStateRef = useRef(null) + const pendingInterruptActionRef = useRef<(() => void) | null>(null) + const agentSessionIdRef = useRef(null) + const conversationItemsRef = useRef([]) + const pendingConfirmIdsRef = useRef([]) + const hasActiveAgentTask = agentStreaming || pendingConfirmIds.length > 0 + + const getDrawerMaxWidth = useCallback(() => { + if (typeof window === 'undefined') return CHAT_DRAWER_MAX_WIDTH + return Math.min(CHAT_DRAWER_MAX_WIDTH, Math.max(CHAT_DRAWER_MIN_WIDTH, window.innerWidth - 48)) + }, []) + + const clampDrawerWidth = useCallback( + (width: number) => Math.min(getDrawerMaxWidth(), Math.max(CHAT_DRAWER_MIN_WIDTH, width)), + [getDrawerMaxWidth] + ) + + const handleDrawerResizeStart = useCallback( + (event: React.PointerEvent) => { + if (event.pointerType === 'mouse' && event.button !== 0) return + event.preventDefault() + event.stopPropagation() + + const handlePointerMove = (moveEvent: PointerEvent) => { + setDrawerWidth(clampDrawerWidth(window.innerWidth - moveEvent.clientX)) + } + const handlePointerUp = () => { + window.removeEventListener('pointermove', handlePointerMove) + window.removeEventListener('pointerup', handlePointerUp) + } + + window.addEventListener('pointermove', handlePointerMove) + window.addEventListener('pointerup', handlePointerUp) + }, + [clampDrawerWidth] + ) + + const copyMessageText = useCallback( + async (text: string | undefined) => { + const content = (text ?? '').trim() + if (!content) return + try { + await navigator.clipboard.writeText(content) + toast.success(t('aiops.chat.copied', { defaultValue: '已复制' })) + } catch (error) { + toast.error( + t('aiops.chat.copyFailed', { + defaultValue: '复制失败：{{message}}', + message: error instanceof Error ? error.message : String(error), + }) + ) + } + }, + [t] + ) + + const showAgentServiceError = useCallback( + (error: Error) => { + toast.error( + t('aiops.agent.serviceUnavailable', { + defaultValue: 'AI 服务暂不可用：{{message}}', + message: + error.message || + t('aiops.common.unknownError', { + defaultValue: '未知错误', + }), + }) + ) + }, + [t] + ) + + const handleFeedbackChange = useCallback((fb: AgentFeedback) => { + setFeedbackMap((prev) => ({ ...prev, [`${fb.targetType}:${fb.targetId}`]: fb })) + }, []) + + const { data: agentSessions = [], refetch: refetchAgentSessions } = useQuery({ + queryKey: ['agent-sessions', agentSurface], + queryFn: async () => (await apiListSessions(agentSurface)).data, + enabled: isOpen, + }) + const { data: agentConfigSummary } = useQuery({ + queryKey: ['agent-config-summary'], + queryFn: async () => (await apiGetAgentConfigSummary()).data, + enabled: isOpen, + }) + + const getPageContext = useCallback(() => { + const pathname = window.location.pathname + const jobMatch = pathname.match(/\/jobs\/detail\/([^/?#]+)/) + const nodeMatch = pathname.match(/\/nodes\/([^/?#]+)/) + return { + route: pathname, + url: pathname, + jobName: jobMatch?.[1] ?? currentJobName, + nodeName: nodeMatch?.[1], + entryPoint: inferAgentEntryPoint(pathname), + surface: agentSurface, + } + }, [agentSurface, currentJobName]) + + const getClientContext = useCallback( + () => ({ + locale: typeof navigator !== 'undefined' ? navigator.language : undefined, + timezone: + typeof Intl !== 'undefined' ? Intl.DateTimeFormat().resolvedOptions().timeZone : undefined, + }), + [] + ) + + const cancelAgentStream = useCallback(() => { + agentAbortRef.current?.abort() + agentAbortRef.current = null + setAgentStreaming(false) + }, []) + + const clearThinkingItems = useCallback(() => { + setConversationItems((prev) => prev.filter((item) => item.kind !== 'thinking')) + }, []) + + const rememberFailedAgentRequest = useCallback((request: AgentPendingRequest) => { + setFailedAgentRequests((prev) => ({ + ...prev, + [request.requestId]: request, + })) + }, []) + + const clearFailedAgentRequest = useCallback((requestId: string | undefined) => { + if (!requestId) return + setFailedAgentRequests((prev) => { + if (!prev[requestId]) return prev + const next = { ...prev } + delete next[requestId] + return next + }) + }, []) + + const updateUserRequestState = useCallback( + ( + requestId: string | undefined, + requestState: ConversationItem['requestState'], + requestError?: string + ) => { + if (!requestId) return + setConversationItems((prev) => + prev.map((item) => + item.kind === 'user' && item.requestId === requestId + ? { + ...item, + requestState, + requestError, + } + : item + ) + ) + }, + [] + ) + + const cancelActiveAgentRequest = useCallback(() => { + const activeRequestId = activeAgentRequestStateRef.current?.requestId + cancelAgentStream() + clearThinkingItems() + updateUserRequestState( + activeRequestId, + 'cancelled', + t('aiops.agent.cancelledByUser', { defaultValue: '已取消' }) + ) + activeAgentRequestStateRef.current = null + setRetryableAgentRequest(null) + }, [cancelAgentStream, clearThinkingItems, t, updateUserRequestState]) + + const resolveConfirmation = useCallback((confirmId: string) => { + setPendingConfirmIds((prev) => prev.filter((id) => id !== confirmId)) + }, []) + + useEffect(() => { + pendingConfirmIdsRef.current = pendingConfirmIds + }, [pendingConfirmIds]) + + const bumpAgentInteractionVersion = useCallback(() => { + agentInteractionVersionRef.current += 1 + }, []) + + const requestInterruptConfirmation = useCallback( + (action: () => void, options?: Partial) => { + if (!hasActiveAgentTask) { + action() + return + } + + pendingInterruptActionRef.current = action + setInterruptConfirmState({ + title: options?.title ?? '中断当前 Agent 执行？', + description: + options?.description ?? + '当前 Agent 仍在执行或等待确认。关闭助手、切换模式或切换会话都会中断这轮思考与工具调用。', + confirmLabel: options?.confirmLabel ?? '中断并继续', + }) + }, + [hasActiveAgentTask] + ) + + const cancelInterruptConfirmation = useCallback(() => { + pendingInterruptActionRef.current = null + setInterruptConfirmState(null) + }, []) + + const confirmInterruptAndContinue = useCallback(() => { + const action = pendingInterruptActionRef.current + pendingInterruptActionRef.current = null + setInterruptConfirmState(null) + cancelAgentStream() + action?.() + }, [cancelAgentStream]) + + const cancelSessionDeleteConfirmation = useCallback(() => { + setSessionDeleteConfirmState(null) + }, []) + + const handleToggleSessionPin = useCallback( + async (session: AgentSession) => { + if (sessionActionLoading || agentHistoryLoading) return + const nextPinned = !session.pinnedAt + setSessionActionLoading({ sessionId: session.sessionId, action: 'pin' }) + try { + await apiPinSession(session.sessionId, nextPinned) + await refetchAgentSessions() + } finally { + setSessionActionLoading((current) => + current?.sessionId === session.sessionId && current.action === 'pin' ? null : current + ) + } + }, + [agentHistoryLoading, refetchAgentSessions, sessionActionLoading] + ) + + const beginRenameSession = useCallback((session: AgentSession) => { + setRenameState({ + sessionId: session.sessionId, + draftTitle: session.title ?? '', + }) + // Focus the input on the next tick after it has rendered. + requestAnimationFrame(() => { + renameInputRef.current?.focus() + renameInputRef.current?.select() + }) + }, []) + + const cancelRenameSession = useCallback(() => { + setRenameState(null) + }, []) + + const commitRenameSession = useCallback( + async (session: AgentSession) => { + if (!renameState || renameState.sessionId !== session.sessionId) return + const next = renameState.draftTitle.trim() + if (next === '' || next === (session.title ?? '')) { + setRenameState(null) + return + } + setSessionActionLoading({ sessionId: session.sessionId, action: 'rename' }) + try { + await apiRenameSession(session.sessionId, next) + await refetchAgentSessions() + setRenameState(null) + } catch (err) { + toast.error( + t('aiops.agent.renameSessionFailed', { + defaultValue: '会话改名失败：{{msg}}', + msg: err instanceof Error ? err.message : String(err), + }) + ) + } finally { + setSessionActionLoading((current) => + current?.sessionId === session.sessionId && current.action === 'rename' ? null : current + ) + } + }, + [refetchAgentSessions, renameState, t] + ) + + const requestDeleteAgentSession = useCallback((session: AgentSession) => { + setSessionDeleteConfirmState({ + sessionId: session.sessionId, + title: session.title || '未命名', + }) + }, []) + + const persistAgentSessionId = useCallback( + (sessionId: string | null) => { + agentSessionIdRef.current = sessionId + setSelectedAgentSessionId(sessionId) + if (typeof window === 'undefined') return + if (sessionId) { + window.localStorage.setItem(agentSessionStorageKey, sessionId) + } else { + window.localStorage.removeItem(agentSessionStorageKey) + } + }, + [agentSessionStorageKey] + ) + + const resetAgentConversation = useCallback(() => { + bumpAgentInteractionVersion() + agentHistoryRequestIdRef.current += 1 + cancelAgentStream() + setAgentHistoryLoading(false) + setConversationItems([]) + setPendingConfirmIds([]) + setAgentHistoryError(null) + setRetryableAgentRequest(null) + setFailedAgentRequests({}) + lastAgentHistorySessionIdRef.current = null + lastLoadedAgentSessionIdRef.current = null + lastAgentRequestRef.current = null + activeAgentRequestStateRef.current = null + pendingInterruptActionRef.current = null + setInterruptConfirmState(null) + persistAgentSessionId(null) + }, [bumpAgentInteractionVersion, cancelAgentStream, persistAgentSessionId]) + + const performDeleteAgentSession = useCallback( + async (sessionId: string) => { + setSessionActionLoading({ sessionId, action: 'delete' }) + try { + await apiDeleteSession(sessionId) + const deletingCurrentSession = + selectedAgentSessionId === sessionId || agentSessionIdRef.current === sessionId + if (deletingCurrentSession) { + resetAgentConversation() + } + await refetchAgentSessions() + } finally { + setSessionActionLoading((current) => + current?.sessionId === sessionId && current.action === 'delete' ? null : current + ) + } + }, + [refetchAgentSessions, resetAgentConversation, selectedAgentSessionId] + ) + + const confirmDeleteAgentSession = useCallback(() => { + const pending = sessionDeleteConfirmState + if (!pending) return + setSessionDeleteConfirmState(null) + + const runDelete = () => { + void performDeleteAgentSession(pending.sessionId) + } + + const deletingCurrentSession = + selectedAgentSessionId === pending.sessionId || + agentSessionIdRef.current === pending.sessionId + if (deletingCurrentSession && hasActiveAgentTask) { + requestInterruptConfirmation(runDelete, { + title: t('aiops.agent.interruptDeleteSessionTitle', { + defaultValue: '删除当前会话会中断当前执行', + }), + description: t('aiops.agent.interruptDeleteSessionDescription', { + defaultValue: + '当前 Agent 仍在执行或等待确认。删除这个会话会立即中断本轮流程，并从历史中移除该会话。', + }), + confirmLabel: t('aiops.agent.interruptDeleteSessionConfirm', { + defaultValue: '中断并删除', + }), + }) + return + } + + runDelete() + }, [ + hasActiveAgentTask, + performDeleteAgentSession, + requestInterruptConfirmation, + selectedAgentSessionId, + sessionDeleteConfirmState, + t, + ]) + + const updateToolCallItem = useCallback( + (toolCallId: string | undefined, updater: (item: ConversationItem) => ConversationItem) => { + if (!toolCallId) return + setConversationItems((prev) => + prev.map((item) => + item.id === toolCallId || item.id === `toolcall-${toolCallId}` ? updater(item) : item + ) + ) + }, + [] + ) + + // ── SSE event handler ───────────────────────────────────────────────────── + + const handleAgentSSEEvent = useCallback( + (event: AgentSSEEvent, thinkingId: string) => { + const eventData: AgentEventPayload = + typeof event.data === 'object' && event.data !== null + ? (event.data as AgentEventPayload) + : {} + + switch (event.event) { + case 'agent_run_started': + break + case 'agent_status': + case 'agent_handoff': { + const summary = + eventData.summary ?? eventData.content ?? eventData.description ?? eventData.title ?? '' + const role = eventData.agentRole ?? 'single_agent' + if ( + role === 'single_agent' && + event.event === 'agent_status' && + eventData.status === 'running' + ) { + setConversationItems((prev) => { + const existing = prev.find((item) => item.id === thinkingId) + if (existing) { + return prev.map((item) => + item.id === thinkingId + ? { ...item, thinkingContent: summary || 'Agent 思考中...' } + : item + ) + } + return [ + ...prev, + { + id: thinkingId, + kind: 'thinking', + thinkingContent: summary || 'Agent 思考中...', + timestamp: new Date(), + }, + ] + }) + } + break + } + + case 'thinking': { + const content: string = + typeof event.data === 'string' ? event.data : (eventData.content ?? '') + setConversationItems((prev) => { + const existing = prev.find((i) => i.id === thinkingId) + if (existing) { + return prev.map((i) => + i.id === thinkingId + ? { ...i, thinkingContent: (i.thinkingContent ?? '') + content } + : i + ) + } + return [ + ...prev, + { + id: thinkingId, + kind: 'thinking', + thinkingContent: content, + timestamp: new Date(), + }, + ] + }) + break + } + + case 'tool_call': + case 'tool_call_started': { + const toolName: string = + eventData.toolName ?? eventData.name ?? eventData.tool ?? 'unknown' + const toolArgs: Record = + eventData.toolArgs ?? eventData.args ?? eventData.arguments ?? {} + const toolCallId = + eventData.toolCallId ?? eventData.id ?? `tool-${toolName}-${Date.now()}` + setConversationItems((prev) => { + const itemId = `toolcall-${toolCallId}` + const exists = prev.some((item) => item.id === itemId) + if (exists) { + return prev.map((item) => + item.id === itemId + ? { ...item, toolName, toolArgs, toolStatus: 'executing' as const } + : item + ) + } + return [ + ...prev, + { + id: itemId, + kind: 'tool_call', + toolName, + toolArgs, + toolStatus: 'executing', + timestamp: new Date(), + }, + ] + }) + break + } + + case 'tool_result': + case 'tool_call_completed': { + const toolName = eventData.toolName ?? eventData.name ?? eventData.tool ?? 'unknown' + const toolCallId = + eventData.toolCallId ?? eventData.id ?? `tool-${toolName}-${Date.now()}` + const result: string = + typeof eventData.resultSummary === 'string' + ? eventData.resultSummary + : typeof eventData.result === 'string' + ? eventData.result + : JSON.stringify(eventData.result ?? event.data ?? '') + const isError: boolean = eventData.isError ?? false + setConversationItems((prev) => { + const itemId = `toolcall-${toolCallId}` + const exists = prev.some((item) => item.id === itemId) + if (!exists) { + return [ + ...prev, + { + id: itemId, + kind: 'tool_call', + toolName, + toolArgs: eventData.toolArgs ?? eventData.args ?? eventData.arguments ?? {}, + toolStatus: isError ? 'error' : 'done', + toolResult: result, + timestamp: new Date(), + }, + ] + } + return prev.map((item) => + item.id === itemId + ? { + ...item, + toolStatus: isError ? ('error' as const) : ('done' as const), + toolResult: result, + } + : item + ) + }) + break + } + + case 'message': + case 'final_answer': { + const text: string = + typeof event.data === 'string' ? event.data : (eventData.content ?? '') + if (!text.trim()) { + clearThinkingItems() + break + } + const isAskResponse = eventData.agentRole === 'ask' || eventData.agentId === 'ask-1' + const askStreamItemId = `ask-stream-${ + eventData.turnId || activeAgentRequestStateRef.current?.requestId || thinkingId + }` + if (isAskResponse && eventData.partial === true) { + if (activeAgentRequestStateRef.current) { + activeAgentRequestStateRef.current = { + ...activeAgentRequestStateRef.current, + hasFinalResponse: true, + } + } + if (eventData.sessionId) { + lastLoadedAgentSessionIdRef.current = eventData.sessionId + persistAgentSessionId(eventData.sessionId) + } + clearThinkingItems() + setConversationItems((prev) => { + const existing = prev.find((item) => item.id === askStreamItemId) + if (existing?.kind === 'message') { + return prev.map((item) => + item.id === askStreamItemId + ? { + ...item, + text: `${item.text ?? ''}${text}`, + timestamp: new Date(), + } + : item + ) + } + return [ + ...prev, + { + id: askStreamItemId, + kind: 'message', + text, + timestamp: new Date(), + }, + ] + }) + break + } + if (activeAgentRequestStateRef.current) { + activeAgentRequestStateRef.current = { + ...activeAgentRequestStateRef.current, + hasFinalResponse: true, + } + } + if (eventData.sessionId) { + lastLoadedAgentSessionIdRef.current = eventData.sessionId + persistAgentSessionId(eventData.sessionId) + } + clearThinkingItems() + setConversationItems((prev) => { + if (isAskResponse && prev.some((item) => item.id === askStreamItemId)) { + return prev.map((item) => + item.id === askStreamItemId + ? { + ...item, + text, + feedbackTargetId: eventData.feedbackTargetId ?? item.feedbackTargetId, + timestamp: new Date(), + } + : item + ) + } + return [ + ...prev, + { + id: `msg-${Date.now()}`, + kind: 'message', + text, + feedbackTargetId: eventData.feedbackTargetId, + timestamp: new Date(), + }, + ] + }) + break + } + + case 'confirmation_required': + case 'tool_call_confirmation_required': { + const confirmId = eventData.confirmId ?? eventData.confirm_id ?? eventData.id ?? '' + const confirmToolCallId = eventData.toolCallId + clearThinkingItems() + if (activeAgentRequestStateRef.current) { + activeAgentRequestStateRef.current = { + ...activeAgentRequestStateRef.current, + awaitingConfirmation: true, + } + } + if (confirmId) { + setPendingConfirmIds((prev) => (prev.includes(confirmId) ? prev : [...prev, confirmId])) + } + updateToolCallItem(confirmToolCallId, (item) => ({ + ...item, + toolStatus: 'awaiting_confirmation', + toolResult: eventData.description ?? '等待用户确认后继续执行', + })) + const nextItem: ConversationItem = { + id: `confirm-${Date.now()}`, + kind: 'confirmation_required', + confirmId, + confirmToolCallId, + confirmAction: eventData.action ?? eventData.toolName ?? eventData.tool_name ?? '', + confirmDescription: eventData.description ?? '', + confirmRiskLevel: eventData.riskLevel ?? eventData.risk_level, + confirmPermissionExplanation: + eventData.permissionExplanation ?? eventData.permission_explanation, + confirmRiskExplanation: eventData.riskExplanation ?? eventData.risk_explanation, + confirmAffectedResources: eventData.affectedResources ?? eventData.affected_resources, + confirmInteraction: eventData.interaction ?? 'approval', + confirmForm: eventData.form, + timestamp: new Date(), + } + setConversationItems((prev) => { + const existing = confirmId + ? prev.find( + (item) => item.kind === 'confirmation_required' && item.confirmId === confirmId + ) + : undefined + if (!existing) { + return [...prev, nextItem] + } + return prev.map((item) => + item.id === existing.id + ? { + ...item, + confirmToolCallId, + confirmAction: nextItem.confirmAction, + confirmDescription: nextItem.confirmDescription, + confirmRiskLevel: nextItem.confirmRiskLevel, + confirmPermissionExplanation: nextItem.confirmPermissionExplanation, + confirmRiskExplanation: nextItem.confirmRiskExplanation, + confirmAffectedResources: nextItem.confirmAffectedResources, + confirmInteraction: nextItem.confirmInteraction, + confirmForm: nextItem.confirmForm, + } + : item + ) + }) + break + } + + case 'parameter_review': { + const payload = eventData as unknown as ParameterReviewPayload + clearThinkingItems() + setConversationItems((prev) => [ + ...prev, + { + id: `param-review-${payload.reviewId || Date.now()}`, + kind: 'parameter_review', + parameterReview: payload, + timestamp: new Date(), + }, + ]) + break + } + + default: + break + } + }, + [clearThinkingItems, persistAgentSessionId, updateToolCallItem] + ) + + const loadAgentSession = useCallback( + async (sessionId: string) => { + const requestId = agentHistoryRequestIdRef.current + 1 + const interactionVersion = agentInteractionVersionRef.current + agentHistoryRequestIdRef.current = requestId + lastAgentHistorySessionIdRef.current = sessionId + setSelectedAgentSessionId(sessionId) + setAgentHistoryLoading(true) + setAgentHistoryError(null) + setRetryableAgentRequest(null) + setFailedAgentRequests({}) + setFeedbackMap({}) + cancelAgentStream() + setPendingConfirmIds([]) + activeAgentRequestStateRef.current = null + try { + const [messages, toolCalls, turns, feedbacks] = await Promise.all([ + apiGetSessionMessages(sessionId).then((response) => response.data), + apiGetSessionToolCalls(sessionId).then((response) => response.data), + apiGetSessionTurns(sessionId).then((response) => response.data), + apiListFeedbacks(sessionId) + .then((response) => response.data ?? []) + .catch((): AgentFeedback[] => []), + ]) + const runEventsByTurn = Object.fromEntries( + await Promise.all( + turns.map(async (turn) => [ + turn.turnId, + await apiGetTurnEvents(turn.turnId).then((response) => response.data), + ]) + ) + ) as Record + if ( + requestId !== agentHistoryRequestIdRef.current || + interactionVersion !== agentInteractionVersionRef.current + ) { + return + } + const items = mapSessionHistoryToConversationItems( + messages, + toolCalls, + turns, + runEventsByTurn + ) + setConversationItems(items) + // Populate feedback map keyed by "targetType:targetId" + const fbMap: Record = {} + for (const fb of feedbacks) { + fbMap[`${fb.targetType}:${fb.targetId}`] = fb + } + setFeedbackMap(fbMap) + setPendingConfirmIds( + items + .filter((item) => item.kind === 'confirmation_required' && item.confirmId) + .map((item) => item.confirmId as string) + ) + const session = agentSessions.find((entry) => entry.sessionId === sessionId) + if (session?.lastOrchestrationMode) { + setOrchestrationMode(normalizeAgentOrchestrationMode(session.lastOrchestrationMode)) + } + lastLoadedAgentSessionIdRef.current = sessionId + persistAgentSessionId(sessionId) + } catch (error) { + if (requestId !== agentHistoryRequestIdRef.current) { + return + } + const message = + error instanceof Error + ? error.message + : t('aiops.common.unknownError', { defaultValue: '未知错误' }) + setAgentHistoryError(message) + } finally { + if (requestId === agentHistoryRequestIdRef.current) { + setAgentHistoryLoading(false) + } + } + }, + [agentSessions, cancelAgentStream, persistAgentSessionId, t] + ) + + const retryLoadAgentSession = useCallback(() => { + const sessionId = lastAgentHistorySessionIdRef.current + if (!sessionId || agentHistoryLoading || agentStreaming) return + void loadAgentSession(sessionId) + }, [agentHistoryLoading, agentStreaming, loadAgentSession]) + + useEffect(() => { + if (!isOpen) cancelAgentStream() + }, [cancelAgentStream, isOpen]) + + useEffect(() => { + if (!hasActiveAgentTask || typeof window === 'undefined') return + + const handleBeforeUnload = (event: BeforeUnloadEvent) => { + event.preventDefault() + event.returnValue = '' + } + + window.addEventListener('beforeunload', handleBeforeUnload) + return () => window.removeEventListener('beforeunload', handleBeforeUnload) + }, [hasActiveAgentTask]) + + useEffect(() => { + if (typeof window === 'undefined') return + bumpAgentInteractionVersion() + agentHistoryRequestIdRef.current += 1 + cancelAgentStream() + setAgentHistoryLoading(false) + setConversationItems([]) + setPendingConfirmIds([]) + setAgentHistoryError(null) + setRetryableAgentRequest(null) + setFailedAgentRequests({}) + setFeedbackMap({}) + lastAgentHistorySessionIdRef.current = null + lastLoadedAgentSessionIdRef.current = null + lastAgentRequestRef.current = null + activeAgentRequestStateRef.current = null + pendingInterruptActionRef.current = null + setInterruptConfirmState(null) + const storedSessionId = window.localStorage.getItem(agentSessionStorageKey) + agentSessionIdRef.current = storedSessionId + setSelectedAgentSessionId(storedSessionId) + }, [agentSessionStorageKey, bumpAgentInteractionVersion, cancelAgentStream]) + + useEffect(() => { + if (!isOpen || typeof window === 'undefined') return + if (agentSessionIdRef.current) return + const storedSessionId = window.localStorage.getItem(agentSessionStorageKey) + if (storedSessionId) { + persistAgentSessionId(storedSessionId) + } + }, [agentSessionStorageKey, isOpen, persistAgentSessionId]) + + useEffect(() => { + if (selectedAgentSessionId) return + if (agentConfigSummary?.defaultOrchestrationMode) { + setOrchestrationMode( + normalizeAgentOrchestrationMode(agentConfigSummary.defaultOrchestrationMode) + ) + } + }, [agentConfigSummary?.defaultOrchestrationMode, selectedAgentSessionId]) + + useEffect(() => { + if (!isOpen || conversationItems.length > 0 || agentHistoryLoading) return + const sessionId = agentSessionIdRef.current + if (!sessionId || lastLoadedAgentSessionIdRef.current === sessionId) return + if (!agentSessions.some((session) => session.sessionId === sessionId)) return + void loadAgentSession(sessionId) + }, [agentHistoryLoading, conversationItems.length, agentSessions, isOpen, loadAgentSession]) + + const startAskRequest = useCallback( + (request: AgentPendingRequest, options?: { appendUserBubble?: boolean }) => { + if (agentStreaming || pendingConfirmIds.length > 0) return + + const appendUserBubble = options?.appendUserBubble ?? true + bumpAgentInteractionVersion() + agentHistoryRequestIdRef.current += 1 + setAgentHistoryLoading(false) + setAgentHistoryError(null) + setRetryableAgentRequest(null) + setFailedAgentRequests({}) + activeAgentRequestStateRef.current = { + requestId: request.requestId, + hasFinalResponse: false, + awaitingConfirmation: false, + } + + if (appendUserBubble) { + setConversationItems((prev) => [ + ...prev, + { + id: `user-${Date.now()}`, + kind: 'user', + text: request.message, + requestId: request.requestId, + requestState: 'running', + requestError: undefined, + requestSessionId: request.sessionId, + requestOrchestrationMode: 'ask', + timestamp: new Date(), + }, + ]) + } + + lastAgentRequestRef.current = request + setAgentStreaming(true) + const thinkingId = `thinking-ask-${Date.now()}` + + const ctrl = connectAgentAskStream( + agentSessionIdRef.current ?? request.sessionId, + request.requestId, + request.message, + request.pageContext, + request.clientContext, + request.pageContext.jobName, + (event: AgentSSEEvent) => handleAgentSSEEvent(event, thinkingId), + (err: Error) => { + agentAbortRef.current = null + clearThinkingItems() + showAgentServiceError(err) + updateUserRequestState(request.requestId, 'failed', err.message) + rememberFailedAgentRequest(request) + setRetryableAgentRequest(request) + activeAgentRequestStateRef.current = null + setAgentStreaming(false) + }, + () => { + agentAbortRef.current = null + clearThinkingItems() + const activeState = activeAgentRequestStateRef.current + const isCurrentRequest = activeState?.requestId === request.requestId + if (isCurrentRequest && activeState?.hasFinalResponse) { + updateUserRequestState(request.requestId, 'done', undefined) + clearFailedAgentRequest(request.requestId) + setRetryableAgentRequest(null) + } else { + updateUserRequestState( + request.requestId, + 'failed', + t('aiops.agent.missingFinalAnswer', { + defaultValue: getFailedRequestMessage(), + }) + ) + rememberFailedAgentRequest(request) + setRetryableAgentRequest(request) + } + activeAgentRequestStateRef.current = null + setAgentStreaming(false) + void refetchAgentSessions() + }, + (sessionId: string) => { + lastLoadedAgentSessionIdRef.current = sessionId + persistAgentSessionId(sessionId) + if (lastAgentRequestRef.current?.requestId === request.requestId) { + lastAgentRequestRef.current = { + ...lastAgentRequestRef.current, + sessionId, + } + } + } + ) + + agentAbortRef.current = ctrl + }, + [ + agentStreaming, + pendingConfirmIds.length, + bumpAgentInteractionVersion, + clearFailedAgentRequest, + clearThinkingItems, + handleAgentSSEEvent, + persistAgentSessionId, + refetchAgentSessions, + rememberFailedAgentRequest, + showAgentServiceError, + t, + updateUserRequestState, + ] + ) + + const startAgentRequest = useCallback( + (request: AgentPendingRequest, options?: { appendUserBubble?: boolean }) => { + if (agentStreaming || pendingConfirmIds.length > 0) return + + bumpAgentInteractionVersion() + agentHistoryRequestIdRef.current += 1 + const appendUserBubble = options?.appendUserBubble ?? true + setAgentHistoryLoading(false) + setAgentHistoryError(null) + setRetryableAgentRequest(null) + clearFailedAgentRequest(request.requestId) + activeAgentRequestStateRef.current = { + requestId: request.requestId, + hasFinalResponse: false, + awaitingConfirmation: false, + } + + if (appendUserBubble) { + setConversationItems((prev) => [ + ...prev, + { + id: `user-${Date.now()}`, + kind: 'user', + text: request.message, + requestId: request.requestId, + requestState: 'running', + requestError: undefined, + requestSessionId: request.sessionId, + requestOrchestrationMode: request.orchestrationMode, + timestamp: new Date(), + }, + ]) + } + + lastAgentRequestRef.current = request + setAgentStreaming(true) + + const thinkingId = `thinking-${Date.now()}` + + const ctrl = connectAgentChat( + agentSessionIdRef.current ?? request.sessionId, + request.requestId, + request.message, + request.pageContext, + request.orchestrationMode, + request.clientContext, + (event: AgentSSEEvent) => handleAgentSSEEvent(event, thinkingId), + (err: Error) => { + agentAbortRef.current = null + clearThinkingItems() + showAgentServiceError(err) + updateUserRequestState(request.requestId, 'failed', err.message) + rememberFailedAgentRequest(request) + setRetryableAgentRequest(request) + activeAgentRequestStateRef.current = null + setAgentStreaming(false) + }, + () => { + agentAbortRef.current = null + clearThinkingItems() + const activeState = activeAgentRequestStateRef.current + const isCurrentRequest = activeState?.requestId === request.requestId + const hasFinalResponse = isCurrentRequest && activeState?.hasFinalResponse + const awaitingConfirmation = isCurrentRequest && activeState?.awaitingConfirmation + + if (hasFinalResponse) { + updateUserRequestState(request.requestId, 'done', undefined) + clearFailedAgentRequest(request.requestId) + setRetryableAgentRequest(null) + } else if (awaitingConfirmation) { + updateUserRequestState(request.requestId, 'awaiting_confirmation', undefined) + clearFailedAgentRequest(request.requestId) + setRetryableAgentRequest(null) + } else { + updateUserRequestState( + request.requestId, + 'failed', + t('aiops.agent.missingFinalAnswer', { + defaultValue: getFailedRequestMessage(), + }) + ) + rememberFailedAgentRequest(request) + setRetryableAgentRequest(request) + } + activeAgentRequestStateRef.current = null + setAgentStreaming(false) + void refetchAgentSessions() + }, + (sessionId: string) => { + lastLoadedAgentSessionIdRef.current = sessionId + persistAgentSessionId(sessionId) + if (lastAgentRequestRef.current?.requestId === request.requestId) { + lastAgentRequestRef.current = { + ...lastAgentRequestRef.current, + sessionId, + } + } + } + ) + + agentAbortRef.current = ctrl + }, + [ + agentStreaming, + pendingConfirmIds.length, + bumpAgentInteractionVersion, + clearFailedAgentRequest, + persistAgentSessionId, + rememberFailedAgentRequest, + refetchAgentSessions, + clearThinkingItems, + handleAgentSSEEvent, + showAgentServiceError, + t, + updateUserRequestState, + ] + ) + + const startAgentResume = useCallback( + (confirmId: string) => { + if (!confirmId || agentStreaming) { + return + } + + bumpAgentInteractionVersion() + agentHistoryRequestIdRef.current += 1 + setAgentHistoryLoading(false) + setAgentHistoryError(null) + setRetryableAgentRequest(null) + setAgentStreaming(true) + const thinkingId = `thinking-resume-${Date.now()}` + + const ctrl = connectAgentResume( + confirmId, + (event: AgentSSEEvent) => handleAgentSSEEvent(event, thinkingId), + (err: Error) => { + agentAbortRef.current = null + clearThinkingItems() + setConversationItems((prev) => [ + ...prev, + { + id: `confirm-error-${Date.now()}`, + kind: 'error', + text: err.message, + timestamp: new Date(), + }, + ]) + setAgentStreaming(false) + }, + () => { + agentAbortRef.current = null + clearThinkingItems() + setAgentStreaming(false) + void refetchAgentSessions() + }, + (sessionId: string) => { + lastLoadedAgentSessionIdRef.current = sessionId + persistAgentSessionId(sessionId) + } + ) + + agentAbortRef.current = ctrl + }, + [ + agentStreaming, + bumpAgentInteractionVersion, + clearThinkingItems, + handleAgentSSEEvent, + persistAgentSessionId, + refetchAgentSessions, + ] + ) + + const handleAgentSend = (messageText?: string) => { + const textToSend = messageText ?? input.trim() + if (!textToSend || agentStreaming || pendingConfirmIds.length > 0) return + if (textToSend.length > AGENT_INPUT_MAX_LENGTH) { + setConversationItems((prev) => [ + ...prev, + { + id: `err-${Date.now()}`, + kind: 'error', + text: `输入内容超过 ${AGENT_INPUT_MAX_LENGTH} 字，请精简后再发送。`, + timestamp: new Date(), + }, + ]) + return + } + + setInput('') + const request: AgentPendingRequest = { + requestId: generateAgentRequestId(), + sessionId: agentSessionIdRef.current, + message: textToSend, + orchestrationMode: normalizeAgentOrchestrationMode(orchestrationMode), + pageContext: getPageContext(), + clientContext: getClientContext(), + } + if (chatMode === 'llm') { + startAskRequest(request, { appendUserBubble: true }) + return + } + startAgentRequest(request, { appendUserBubble: true }) + } + + const beginEditAgentMessage = useCallback((item: ConversationItem) => { + setEditingAgentMessageId(item.id) + setEditingAgentDraft(item.text ?? '') + }, []) + + const cancelEditAgentMessage = useCallback(() => { + setEditingAgentMessageId(null) + setEditingAgentDraft('') + }, []) + + const submitEditAgentMessage = useCallback( + (item: ConversationItem) => { + const nextContent = editingAgentDraft.trim() + if (!nextContent || agentStreaming || pendingConfirmIds.length > 0) return + + const requestId = generateAgentRequestId() + const nextOrchestrationMode = normalizeAgentOrchestrationMode( + item.requestOrchestrationMode ?? orchestrationMode + ) + setPendingConfirmIds([]) + setRetryableAgentRequest(null) + setFailedAgentRequests({}) + activeAgentRequestStateRef.current = null + lastAgentRequestRef.current = null + + setConversationItems((prev) => { + const index = prev.findIndex((entry) => entry.id === item.id) + if (index < 0) return prev + return prev.slice(0, index + 1).map((entry, currentIndex) => + currentIndex === index + ? { + ...entry, + text: nextContent, + requestId, + requestState: 'running', + requestError: undefined, + requestSessionId: null, + requestOrchestrationMode: chatMode === 'llm' ? 'ask' : nextOrchestrationMode, + timestamp: new Date(), + } + : entry + ) + }) + setEditingAgentMessageId(null) + setEditingAgentDraft('') + const request: AgentPendingRequest = { + requestId, + sessionId: null, + message: nextContent, + orchestrationMode: nextOrchestrationMode, + pageContext: getPageContext(), + clientContext: getClientContext(), + } + if (chatMode === 'llm') { + startAskRequest(request, { appendUserBubble: false }) + } else { + startAgentRequest(request, { appendUserBubble: false }) + } + }, + [ + agentStreaming, + chatMode, + editingAgentDraft, + getClientContext, + getPageContext, + orchestrationMode, + pendingConfirmIds.length, + startAskRequest, + startAgentRequest, + ] + ) + + const regenerateAgentFromUser = useCallback( + (item: ConversationItem) => { + const content = item.text?.trim() + if (!content || agentStreaming || pendingConfirmIds.length > 0) return + + const requestId = generateAgentRequestId() + const nextOrchestrationMode = normalizeAgentOrchestrationMode( + item.requestOrchestrationMode ?? orchestrationMode + ) + setPendingConfirmIds([]) + setRetryableAgentRequest(null) + setFailedAgentRequests({}) + activeAgentRequestStateRef.current = null + lastAgentRequestRef.current = null + + setConversationItems((prev) => { + const index = prev.findIndex((entry) => entry.id === item.id) + if (index < 0) return prev + return prev.slice(0, index + 1).map((entry, currentIndex) => + currentIndex === index + ? { + ...entry, + requestId, + requestState: 'running', + requestError: undefined, + requestSessionId: null, + requestOrchestrationMode: chatMode === 'llm' ? 'ask' : nextOrchestrationMode, + timestamp: new Date(), + } + : entry + ) + }) + + const request: AgentPendingRequest = { + requestId, + sessionId: null, + message: content, + orchestrationMode: nextOrchestrationMode, + pageContext: getPageContext(), + clientContext: getClientContext(), + } + if (chatMode === 'llm') { + startAskRequest(request, { appendUserBubble: false }) + } else { + startAgentRequest(request, { appendUserBubble: false }) + } + }, + [ + agentStreaming, + chatMode, + getClientContext, + getPageContext, + orchestrationMode, + pendingConfirmIds.length, + startAskRequest, + startAgentRequest, + ] + ) + + const retryAgentRequest = useCallback( + (source?: string | ConversationItem) => { + const requestId = typeof source === 'string' ? source : source?.requestId + const requestFromMessage = + source && typeof source !== 'string' && source.kind === 'user' && source.text + ? { + requestId: source.requestId ?? generateAgentRequestId(), + sessionId: agentSessionIdRef.current ?? source.requestSessionId ?? null, + message: source.text, + orchestrationMode: normalizeAgentOrchestrationMode( + source.requestOrchestrationMode ?? orchestrationMode + ), + pageContext: getPageContext(), + clientContext: getClientContext(), + } + : null + const request = + requestFromMessage ?? + (requestId ? failedAgentRequests[requestId] : undefined) ?? + retryableAgentRequest ?? + lastAgentRequestRef.current + if (!request || agentStreaming || pendingConfirmIds.length > 0) return + clearFailedAgentRequest(requestId ?? request.requestId) + const retryRequest: AgentPendingRequest = { + ...request, + requestId: generateAgentRequestId(), + sessionId: agentSessionIdRef.current ?? request.sessionId, + orchestrationMode: normalizeAgentOrchestrationMode(request.orchestrationMode), + pageContext: getPageContext(), + clientContext: getClientContext(), + } + if (chatMode === 'llm') { + startAskRequest(retryRequest, { appendUserBubble: true }) + } else { + startAgentRequest(retryRequest, { appendUserBubble: true }) + } + }, + [ + agentStreaming, + chatMode, + clearFailedAgentRequest, + failedAgentRequests, + getClientContext, + getPageContext, + orchestrationMode, + pendingConfirmIds.length, + retryableAgentRequest, + startAskRequest, + startAgentRequest, + ] + ) + + const retryAgentRequestInNewSession = useCallback( + (source?: string | ConversationItem) => { + const requestId = typeof source === 'string' ? source : source?.requestId + const requestFromMessage = + source && typeof source !== 'string' && source.kind === 'user' && source.text + ? { + requestId: source.requestId ?? generateAgentRequestId(), + sessionId: null, + message: source.text, + orchestrationMode: normalizeAgentOrchestrationMode( + source.requestOrchestrationMode ?? orchestrationMode + ), + pageContext: getPageContext(), + clientContext: getClientContext(), + } + : null + const request = + requestFromMessage ?? + (requestId ? failedAgentRequests[requestId] : undefined) ?? + retryableAgentRequest ?? + lastAgentRequestRef.current + if (!request || agentStreaming || pendingConfirmIds.length > 0) return + clearFailedAgentRequest(requestId ?? request.requestId) + resetAgentConversation() + const retryRequest: AgentPendingRequest = { + ...request, + requestId: generateAgentRequestId(), + sessionId: null, + orchestrationMode: normalizeAgentOrchestrationMode(request.orchestrationMode), + pageContext: getPageContext(), + clientContext: getClientContext(), + } + if (chatMode === 'llm') { + startAskRequest(retryRequest, { appendUserBubble: true }) + } else { + startAgentRequest(retryRequest, { appendUserBubble: true }) + } + }, + [ + agentStreaming, + chatMode, + clearFailedAgentRequest, + failedAgentRequests, + getClientContext, + getPageContext, + orchestrationMode, + pendingConfirmIds.length, + resetAgentConversation, + retryableAgentRequest, + startAskRequest, + startAgentRequest, + ] + ) + + const invalidateAgentAffectedQueries = useCallback( + (toolName: string, status: string) => { + const normalizedTool = String(toolName || '').trim() + const normalizedStatus = String(status || '') + .trim() + .toLowerCase() + + if (!normalizedTool) { + return + } + + void queryClient.invalidateQueries({ queryKey: ['operation-logs'] }) + + if (normalizedStatus === 'rejected') { + return + } + + const jobAffectingTools = new Set([ + 'create_jupyter_job', + 'create_webide_job', + 'create_pytorch_job', + 'create_tensorflow_job', + 'create_custom_job', + 'resubmit_job', + 'stop_job', + 'delete_job', + ]) + if (jobAffectingTools.has(normalizedTool)) { + void queryClient.invalidateQueries({ queryKey: ['vcjobs'] }) + } + }, + [queryClient] + ) + + const handleConfirmationSettled = useCallback( + ( + item: ConversationItem, + result: { status: string; result?: unknown; message?: string }, + nextStatus: ConversationItem['toolStatus'], + fallbackText: string + ) => { + resolveConfirmation(item.confirmId ?? '') + const nextToolResult = + typeof result.result === 'string' + ? result.result + : JSON.stringify(result.result ?? result.message ?? fallbackText) + updateToolCallItem(item.confirmToolCallId, (toolItem) => ({ + ...toolItem, + toolStatus: nextStatus, + toolResult: nextToolResult, + })) + invalidateAgentAffectedQueries(item.confirmAction ?? '', result.status) + const normalizedStatus = String(result.status || '') + .trim() + .toLowerCase() + const remainingConfirmIds = pendingConfirmIdsRef.current.filter( + (id) => id !== (item.confirmId ?? '') + ) + if (remainingConfirmIds.length === 0) { + const rejected = normalizedStatus === 'rejected' + const nextRequestState: ConversationItem['requestState'] = rejected ? 'cancelled' : 'done' + const rejectedMessage: ConversationItem[] = rejected + ? [ + { + id: `confirm-rejected-${Date.now()}`, + kind: 'message', + text: t('aiops.agent.operationRejectedMessage', { + defaultValue: + '已取消该操作，我不会继续执行或再次确认。需要其他操作时请重新告诉我。', + }), + timestamp: new Date(), + }, + ] + : [] + setConversationItems((prev) => + prev + .map((ci) => + ci.kind === 'user' && ci.requestState === 'awaiting_confirmation' + ? { + ...ci, + requestState: nextRequestState, + requestError: rejected + ? t('aiops.agent.operationRejected', { + defaultValue: '操作已取消', + }) + : undefined, + } + : ci + ) + .concat(rejectedMessage) + ) + if (rejected) { + activeAgentRequestStateRef.current = null + setRetryableAgentRequest(null) + return + } + startAgentResume(item.confirmId ?? '') + } + }, + [invalidateAgentAffectedQueries, resolveConfirmation, startAgentResume, t, updateToolCallItem] + ) + + useEffect(() => { + conversationItemsRef.current = conversationItems + }, [conversationItems]) + + useEffect(() => { + messagesEndRef.current?.scrollIntoView({ behavior: 'smooth' }) + }, [conversationItems]) + + useEffect(() => { + resizeChatTextarea(agentInputRef.current) + }, [chatMode, input, isOpen]) + + const handleAgentDrawerClose = useCallback(() => { + requestInterruptConfirmation( + () => { + cancelAgentStream() + onClose() + }, + { + title: t('aiops.agent.interruptCloseTitle', { + defaultValue: '关闭助手会中断当前执行', + }), + description: t('aiops.agent.interruptCloseDescription', { + defaultValue: + '当前 Agent 还在思考或执行工具。关闭助手会立刻中断这轮执行，未完成的答复不会保留。', + }), + confirmLabel: t('aiops.agent.interruptCloseConfirm', { + defaultValue: '中断并关闭', + }), + } + ) + }, [cancelAgentStream, onClose, requestInterruptConfirmation, t]) + + const handleAgentModeSwitch = useCallback( + (nextMode: ChatMode) => { + if (nextMode === chatMode) return + if (!hasActiveAgentTask) { + setChatMode(nextMode) + return + } + + requestInterruptConfirmation( + () => { + cancelAgentStream() + setChatMode(nextMode) + }, + { + title: t('aiops.agent.interruptModeTitle', { + defaultValue: '切换模式会中断当前执行', + }), + description: t('aiops.agent.interruptModeDescription', { + defaultValue: '当前助手还在回答、执行或等待确认。切换模式会立刻中断这一轮流程。', + }), + confirmLabel: t('aiops.agent.interruptModeConfirm', { + defaultValue: '中断并切换', + }), + } + ) + }, + [cancelAgentStream, chatMode, hasActiveAgentTask, requestInterruptConfirmation, t] + ) + + const handleSelectAgentSession = useCallback( + (sessionId: string) => { + if (sessionId === selectedAgentSessionId && !agentHistoryLoading) return + requestInterruptConfirmation( + () => { + void loadAgentSession(sessionId) + }, + { + title: t('aiops.agent.interruptSessionTitle', { + defaultValue: '切换会话会中断当前执行', + }), + description: t('aiops.agent.interruptSessionDescription', { + defaultValue: + '当前 Agent 还在执行或等待确认。切换到其他历史会话会中断当前这一轮执行并载入新的会话内容。', + }), + confirmLabel: t('aiops.agent.interruptSessionConfirm', { + defaultValue: '中断并切换', + }), + } + ) + }, + [agentHistoryLoading, loadAgentSession, requestInterruptConfirmation, selectedAgentSessionId, t] + ) + + const handleCreateAgentSession = useCallback(() => { + requestInterruptConfirmation( + () => { + resetAgentConversation() + }, + { + title: t('aiops.agent.interruptNewSessionTitle', { + defaultValue: '新会话会中断当前执行', + }), + description: t('aiops.agent.interruptNewSessionDescription', { + defaultValue: '当前 Agent 还在执行或等待确认。创建新会话会清空当前上下文并中断这轮执行。', + }), + confirmLabel: t('aiops.agent.interruptNewSessionConfirm', { + defaultValue: '中断并新建', + }), + } + ) + }, [requestInterruptConfirmation, resetAgentConversation, t]) + + if (!isOpen) return null + + // ── Unified ask/agent chat layout ───────────────────────────────────────── + + return ( +
    { + if (event.target === event.currentTarget) { + handleAgentDrawerClose() + } + }} + > +
    event.stopPropagation()} + > +
    + +
    + { + if (!open) { + cancelInterruptConfirmation() + } + }} + > + + + + {interruptConfirmState?.title ?? '中断当前 Agent 执行？'} + + + {interruptConfirmState?.description ?? + '当前 Agent 仍在执行或等待确认，继续操作会中断这轮流程。'} + + + + {t('common.cancel', { defaultValue: '取消' })} + + {interruptConfirmState?.confirmLabel ?? '中断并继续'} + + + + + { + if (!open) { + cancelSessionDeleteConfirmation() + } + }} + > + + + + {t('aiops.agent.deleteSessionTitle', { + defaultValue: '确认删除会话', + })} + + + {t('aiops.agent.deleteSessionDescription', { + defaultValue: '确认要删除「{{title}}」会话吗？删除后无法找回。', + title: sessionDeleteConfirmState?.title ?? '未命名', + })} + + + + {t('common.cancel', { defaultValue: '取消' })} + + {t('aiops.agent.deleteSessionConfirm', { + defaultValue: '删除会话', + })} + + + + + {/* Header */} +
    +
    + +

    {t('aiops.chat.assistantName')}

    +
    +
    + + {chatMode === 'agent' + ? t('aiops.chat.mode.agent', { defaultValue: 'agent' }) + : t('aiops.chat.mode.ask', { defaultValue: 'ask' })} + + + +
    +
    + {showHelp && setShowHelp(false)} />} + {hasActiveAgentTask && !showHelp && ( +
    +

    + {t('aiops.agent.interruptHint', { + defaultValue: + '当前 Agent 正在执行或等待确认。关闭助手、切换模式或切换会话都会中断这轮流程。', + })} +

    +
    + )} + + {/* Agent body: left session panel + right conversation */} +
    + {/* Left: Session list panel — narrow, scrollable */} + {sessionPanelOpen && ( +
    +
    + + {t('aiops.agent.sessionHistory', { defaultValue: '历史' })} + +
    + + +
    +
    + {hasActiveAgentTask && ( +
    +

    + {t('aiops.agent.sessionInterruptNotice', { + defaultValue: '切换会话或新建会话会中断当前执行。', + })} +

    +
    + )} + +
    + {agentSessions.map((session) => { + const isSelected = selectedAgentSessionId === session.sessionId + const isPinning = + sessionActionLoading?.sessionId === session.sessionId && + sessionActionLoading.action === 'pin' + const isDeleting = + sessionActionLoading?.sessionId === session.sessionId && + sessionActionLoading.action === 'delete' + const isRenaming = + sessionActionLoading?.sessionId === session.sessionId && + sessionActionLoading.action === 'rename' + const isEditing = renameState?.sessionId === session.sessionId + const sessionTitle = + session.title || t('aiops.agent.untitledSession', { defaultValue: '未命名' }) + + // Truncate title: show first 10 chars + ellipsis if longer + const displayTitle = + [...sessionTitle].length > 10 + ? [...sessionTitle].slice(0, 10).join('') + '…' + : sessionTitle + + return ( +
    + {isEditing ? ( +
    + + setRenameState({ + sessionId: session.sessionId, + draftTitle: e.target.value, + }) + } + onBlur={() => void commitRenameSession(session)} + onKeyDown={(e) => { + if (e.key === 'Enter') { + e.preventDefault() + void commitRenameSession(session) + } else if (e.key === 'Escape') { + e.preventDefault() + cancelRenameSession() + } + }} + disabled={isRenaming} + /> + {isRenaming && ( + + )} +
    + ) : ( + + )} +
    + + {formatAgentSessionDate(session.updatedAt)} + + + + + + +
    +
    + ) + })} + {agentSessions.length === 0 && ( +
    + {t('aiops.agent.noSessions', { defaultValue: '暂无历史' })} +
    + )} +
    +     + {agentHistoryLoading && ( +

    + {t('aiops.agent.loadingSession', { defaultValue: '加载中…' })} +

    + )} + {agentHistoryError && ( +
    +

    {agentHistoryError}

    + {lastAgentHistorySessionIdRef.current && ( + + )} +
    + )} +
    + )} + + {/* Right: Conversation area */} +
    + {/* Messages */} + +
    + {conversationItems.length === 0 && ( +
    +
    +

    + {chatMode === 'llm' + ? t('aiops.chat.initialMessageAsk', { + defaultValue: + '你好！我是 Crater ask 助手，可以回答平台、作业和排障问题。需要执行创建、停止、删除等操作时，请在底部切换到 agent。', + }) + : t('aiops.agent.initialMessage', { + defaultValue: + '你好！我是 Crater agent，可以自主执行操作来帮助你管理作业。请告诉我你需要什么帮助。', + })} +

    +
    +
    + )} + {conversationItems.map((item, itemIndex) => { + if (item.kind === 'user') { + const canRetry = item.requestState === 'failed' && !!item.text + + return ( +
    + {editingAgentMessageId !== item.id && ( +
    + + + +
    + )} +
    + {editingAgentMessageId === item.id ? ( +
    +