prompt_enqueue.go 3.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899
  1. package securityaudit
  2. import (
  3. "context"
  4. "errors"
  5. )
  6. type Enqueuer struct {
  7. config ConfigStore
  8. repo JobRepository
  9. payload PayloadStore
  10. metrics Metrics
  11. }
  12. func NewEnqueuer(config ConfigStore, repo JobRepository, payload PayloadStore, metrics ...Metrics) *Enqueuer {
  13. var metric Metrics
  14. if len(metrics) > 0 {
  15. metric = metrics[0]
  16. }
  17. return &Enqueuer{config: config, repo: repo, payload: payload, metrics: metric}
  18. }
  19. func (e *Enqueuer) Enqueue(ctx context.Context, req Request) error {
  20. if e == nil || e.config == nil || e.repo == nil || e.payload == nil {
  21. return errors.New("prompt audit enqueuer unavailable")
  22. }
  23. cfg, ok := e.config.Active()
  24. baseFields := requestLogFields(req)
  25. if !ok || cfg.EffectiveMode() != ModeAsync {
  26. LogInfo(EventEnqueueSkipped, mergeLogFields(baseFields, map[string]any{"status": "skipped", "error_code": "mode_not_async"}))
  27. return nil
  28. }
  29. baseFields["config_version"] = cfg.ConfigVersion
  30. if !cfg.IncludesGroup(req.GroupID) {
  31. LogInfo(EventEnqueueSkipped, mergeLogFields(baseFields, map[string]any{"status": "skipped", "error_code": "group_out_of_scope"}))
  32. return nil
  33. }
  34. if len(cfg.EnabledEndpoints()) == 0 {
  35. e.recordDropped()
  36. LogWarn(EventEnqueueDropped, mergeLogFields(baseFields, map[string]any{"status": "dropped", "error_code": "no_enabled_endpoint"}))
  37. return nil
  38. }
  39. snapshot, err := ExtractPromptSnapshot(req)
  40. if errors.Is(err, ErrNoPromptText) {
  41. LogInfo(EventEnqueueSkipped, mergeLogFields(baseFields, map[string]any{"status": "skipped", "error_code": "no_user_text"}))
  42. return nil
  43. }
  44. if err != nil {
  45. e.recordDropped()
  46. LogWarn(EventEnqueueDropped, mergeLogFields(baseFields, map[string]any{"status": "dropped", "error_code": "snapshot_invalid"}))
  47. return nil
  48. }
  49. job, err := e.repo.CreateStagingWithCapacity(ctx, snapshot.Redacted(), cfg.ConfigVersion, 3, cfg.QueueCapacity)
  50. if err != nil {
  51. code := "database_unavailable"
  52. if errors.Is(err, ErrQueueFull) {
  53. code = "queue_full"
  54. }
  55. if errors.Is(err, ErrQueueAdmissionBusy) {
  56. code = "queue_admission_busy"
  57. }
  58. LogWarn(EventEnqueueDropped, mergeLogFields(baseFields, map[string]any{
  59. "queue_capacity": cfg.QueueCapacity, "status": "dropped", "error_code": code,
  60. }))
  61. e.recordDropped()
  62. return err
  63. }
  64. if err := e.payload.Set(ctx, job.ID, snapshot.ScanText, DefaultPayloadTTL); err != nil {
  65. _ = e.repo.MarkStagingFailed(ctx, job.ID, "payload_store_failed", "payload store unavailable")
  66. LogWarn(EventEnqueueDropped, mergeLogFields(baseFields, map[string]any{
  67. "job_id": job.ID, "status": "dropped", "error_code": "payload_store_failed",
  68. }))
  69. e.recordDropped()
  70. return err
  71. }
  72. if err := e.repo.PublishQueued(ctx, job.ID); err != nil {
  73. _ = e.payload.Delete(ctx, job.ID)
  74. _ = e.repo.MarkStagingFailed(ctx, job.ID, "queue_publish_failed", "queue publish failed")
  75. LogWarn(EventEnqueueDropped, mergeLogFields(baseFields, map[string]any{
  76. "job_id": job.ID, "status": "dropped", "error_code": "queue_publish_failed",
  77. }))
  78. e.recordDropped()
  79. return err
  80. }
  81. LogInfo(EventJobEnqueued, mergeLogFields(baseFields, map[string]any{
  82. "job_id": job.ID,
  83. "queue_capacity": cfg.QueueCapacity, "status": "queued",
  84. }))
  85. if e.metrics != nil {
  86. e.metrics.IncEnqueued()
  87. }
  88. return nil
  89. }
  90. func (e *Enqueuer) recordDropped() {
  91. if e != nil && e.metrics != nil {
  92. e.metrics.IncDropped()
  93. }
  94. }