Files
sub2api/backend/internal/securityaudit/prompt_enqueue.go
T
李建琦 6d655c9903
Release / update-version (push) Has been cancelled
Release / build-frontend (push) Has been cancelled
Release / release (push) Has been cancelled
Release / sync-version-file (push) Has been cancelled
CI / shell (push) Canceled after 0s
CI / test (push) Canceled after 0s
CI / frontend (push) Canceled after 0s
CI / golangci-lint (push) Canceled after 0s
Security Scan / backend-security (push) Canceled after 0s
Security Scan / frontend-security (push) Canceled after 0s
Sub2API v1.0 - AI API 网关(二开初始版本,基于上游 Wei-Shaw/sub2api)
2026-08-21 18:30:13 +08:00

100 lines
3.3 KiB
Go

package securityaudit
import (
"context"
"errors"
)
type Enqueuer struct {
config ConfigStore
repo JobRepository
payload PayloadStore
metrics Metrics
}
func NewEnqueuer(config ConfigStore, repo JobRepository, payload PayloadStore, metrics ...Metrics) *Enqueuer {
var metric Metrics
if len(metrics) > 0 {
metric = metrics[0]
}
return &Enqueuer{config: config, repo: repo, payload: payload, metrics: metric}
}
func (e *Enqueuer) Enqueue(ctx context.Context, req Request) error {
if e == nil || e.config == nil || e.repo == nil || e.payload == nil {
return errors.New("prompt audit enqueuer unavailable")
}
cfg, ok := e.config.Active()
baseFields := requestLogFields(req)
if !ok || cfg.EffectiveMode() != ModeAsync {
LogInfo(EventEnqueueSkipped, mergeLogFields(baseFields, map[string]any{"status": "skipped", "error_code": "mode_not_async"}))
return nil
}
baseFields["config_version"] = cfg.ConfigVersion
if !cfg.IncludesGroup(req.GroupID) {
LogInfo(EventEnqueueSkipped, mergeLogFields(baseFields, map[string]any{"status": "skipped", "error_code": "group_out_of_scope"}))
return nil
}
if len(cfg.EnabledEndpoints()) == 0 {
e.recordDropped()
LogWarn(EventEnqueueDropped, mergeLogFields(baseFields, map[string]any{"status": "dropped", "error_code": "no_enabled_endpoint"}))
return nil
}
snapshot, err := ExtractPromptSnapshot(req)
if errors.Is(err, ErrNoPromptText) {
LogInfo(EventEnqueueSkipped, mergeLogFields(baseFields, map[string]any{"status": "skipped", "error_code": "no_user_text"}))
return nil
}
if err != nil {
e.recordDropped()
LogWarn(EventEnqueueDropped, mergeLogFields(baseFields, map[string]any{"status": "dropped", "error_code": "snapshot_invalid"}))
return nil
}
job, err := e.repo.CreateStagingWithCapacity(ctx, snapshot.Redacted(), cfg.ConfigVersion, 3, cfg.QueueCapacity)
if err != nil {
code := "database_unavailable"
if errors.Is(err, ErrQueueFull) {
code = "queue_full"
}
if errors.Is(err, ErrQueueAdmissionBusy) {
code = "queue_admission_busy"
}
LogWarn(EventEnqueueDropped, mergeLogFields(baseFields, map[string]any{
"queue_capacity": cfg.QueueCapacity, "status": "dropped", "error_code": code,
}))
e.recordDropped()
return err
}
if err := e.payload.Set(ctx, job.ID, snapshot.ScanText, DefaultPayloadTTL); err != nil {
_ = e.repo.MarkStagingFailed(ctx, job.ID, "payload_store_failed", "payload store unavailable")
LogWarn(EventEnqueueDropped, mergeLogFields(baseFields, map[string]any{
"job_id": job.ID, "status": "dropped", "error_code": "payload_store_failed",
}))
e.recordDropped()
return err
}
if err := e.repo.PublishQueued(ctx, job.ID); err != nil {
_ = e.payload.Delete(ctx, job.ID)
_ = e.repo.MarkStagingFailed(ctx, job.ID, "queue_publish_failed", "queue publish failed")
LogWarn(EventEnqueueDropped, mergeLogFields(baseFields, map[string]any{
"job_id": job.ID, "status": "dropped", "error_code": "queue_publish_failed",
}))
e.recordDropped()
return err
}
LogInfo(EventJobEnqueued, mergeLogFields(baseFields, map[string]any{
"job_id": job.ID,
"queue_capacity": cfg.QueueCapacity, "status": "queued",
}))
if e.metrics != nil {
e.metrics.IncEnqueued()
}
return nil
}
func (e *Enqueuer) recordDropped() {
if e != nil && e.metrics != nil {
e.metrics.IncDropped()
}
}