From ecfe936b1b4bd8817ab8c8bb8e25f720c9799bf5 Mon Sep 17 00:00:00 2001 From: Vishnuvardhan Reddy Date: Sun, 1 Mar 2026 15:22:37 +0000 Subject: [PATCH] feat(audit): Phase 2 - Core audit logging package Implements comprehensive audit logging infrastructure: - types.go: Structured entry types (Entry, ToolCallData, MessageData, etc.) - logger.go: Async Logger with filtering and channel context support - rotation.go: Size/time-based rotation with gzip compression - global.go: Package-level convenience functions Features: - JSON and text format support - Async write with buffered channel (1000 entries) - Daily rotation + size-based rotation - Automatic cleanup (age, count limits) - Context propagation for request tracing - Nil-safe global logger for convenience - Secure file permissions (0600) The package is production-ready with proper concurrency safety and graceful shutdown support. --- pkg/audit/global.go | 160 ++++++++++++++++++ pkg/audit/logger.go | 380 ++++++++++++++++++++++++++++++++++++++++++ pkg/audit/rotation.go | 276 ++++++++++++++++++++++++++++++ pkg/audit/types.go | 136 +++++++++++++++ 4 files changed, 952 insertions(+) create mode 100644 pkg/audit/global.go create mode 100644 pkg/audit/logger.go create mode 100644 pkg/audit/rotation.go create mode 100644 pkg/audit/types.go diff --git a/pkg/audit/global.go b/pkg/audit/global.go new file mode 100644 index 000000000..889c8b808 --- /dev/null +++ b/pkg/audit/global.go @@ -0,0 +1,160 @@ +package audit + +import ( + "context" + "fmt" + "sync" + "time" + + "github.com/sipeed/picoclaw/pkg/config" +) + +// global provides a package-level default logger for convenience. +// It is safe to use even if not initialized - all methods become no-ops. +var global = &globalLogger{} + +// globalLogger wraps a *Logger with nil safety. +type globalLogger struct { + logger *Logger + mu sync.RWMutex +} + +func (g *globalLogger) set(l *Logger) { + g.mu.Lock() + defer g.mu.Unlock() + g.logger = l +} + +func (g *globalLogger) get() *Logger { + g.mu.RLock() + defer g.mu.RUnlock() + return g.logger +} + +// InitGlobal initializes the global audit logger. +// This should be called once during application startup. +func InitGlobal(cfg config.AuditConfig, workspace string) error { + if !cfg.Enabled { + return nil + } + + logger, err := New(cfg, workspace) + if err != nil { + return err + } + + global.set(logger) + return nil +} + +// CloseGlobal closes the global audit logger. +// This should be called during graceful shutdown. +func CloseGlobal() error { + if l := global.get(); l != nil { + return l.Close() + } + return nil +} + +// Global returns the global logger instance. +// Returns nil if audit logging is disabled. +func Global() *Logger { + return global.get() +} + +// Convenience functions that delegate to the global logger. +// These are no-ops if audit logging is disabled. + +// Log writes a single audit entry to the global logger. +func Log(entry *Entry) { + if l := global.get(); l != nil { + l.Log(entry) + } +} + +// LogToolCall logs a tool execution to the global logger. +func LogToolCall(ctx context.Context, data *ToolCallData, duration int64) { + if l := global.get(); l != nil { + l.Log(&Entry{ + Level: LevelInfo, + Component: "tool", + EventType: EventToolCall, + RequestID: RequestIDFromContext(ctx), + SessionID: SessionIDFromContext(ctx), + AgentID: AgentIDFromContext(ctx), + ToolCall: data, + DurationMs: duration, + }) + } +} + +// LogMessage logs a message event to the global logger. +func LogMessage(ctx context.Context, direction, contentType, content, messageID string) { + if l := global.get(); l != nil { + // Truncate content + const maxLen = 10000 + if len(content) > maxLen { + content = content[:maxLen] + "... [truncated]" + } + + l.Log(&Entry{ + Level: LevelInfo, + Component: "channel", + EventType: EventMessage, + RequestID: RequestIDFromContext(ctx), + SessionID: SessionIDFromContext(ctx), + AgentID: AgentIDFromContext(ctx), + Message: &MessageData{ + Direction: direction, + ContentType: contentType, + Content: content, + MessageID: messageID, + }, + }) + } +} + +// LogError logs an error event to the global logger. +func LogError(ctx context.Context, errorType, message string, recoverable bool) { + if l := global.get(); l != nil { + l.Log(&Entry{ + Level: LevelError, + Component: "system", + EventType: EventError, + RequestID: RequestIDFromContext(ctx), + SessionID: SessionIDFromContext(ctx), + AgentID: AgentIDFromContext(ctx), + Error: &ErrorData{ + ErrorType: errorType, + Message: message, + Recoverable: recoverable, + }, + }) + } +} + +// LogSystem logs a system event to the global logger. +func LogSystem(ctx context.Context, operation string, details map[string]interface{}) { + if l := global.get(); l != nil { + l.Log(&Entry{ + Level: LevelInfo, + Component: "system", + EventType: EventSystem, + RequestID: RequestIDFromContext(ctx), + SessionID: SessionIDFromContext(ctx), + System: &SystemData{ + Operation: operation, + Details: details, + }, + }) + } +} + +// GenerateRequestID generates a unique request ID for context propagation. +// This should be called at the entry point of each request. +func GenerateRequestID() string { + // Simple timestamp-based ID; consider UUID for distributed systems + return fmt.Sprintf("req-%d", time.Now().UnixNano()) +} + + diff --git a/pkg/audit/logger.go b/pkg/audit/logger.go new file mode 100644 index 000000000..9e4cf579b --- /dev/null +++ b/pkg/audit/logger.go @@ -0,0 +1,380 @@ +package audit + +import ( + "context" + "encoding/json" + "fmt" + "os" + "path/filepath" + "sync" + "time" + + "github.com/sipeed/picoclaw/pkg/config" +) + +// Logger provides structured audit logging with rotation and filtering. +// +// The logger is safe for concurrent use. All methods are non-blocking; +// entries are queued for async writing to minimize performance impact. +type Logger struct { + config config.AuditConfig + + // Buffered channel for async writing + entries chan *Entry + + // Writer handles file I/O and rotation + writer *rotatingWriter + + // Worker control + ctx context.Context + cancel context.CancelFunc + wg sync.WaitGroup + + // Event filtering + filter *eventFilter + + // Closed flag for safe shutdown + closed bool + mu sync.RWMutex +} + +// New creates a new audit logger with the given configuration. +// +// The workspace parameter is used to resolve relative paths in the +// audit configuration. If the audit location is already absolute, +// it is used as-is. +// +// Returns nil if audit logging is disabled in the configuration. +func New(cfg config.AuditConfig, workspace string) (*Logger, error) { + if !cfg.Enabled { + return nil, nil + } + + // Resolve log directory + logDir := cfg.Location + if !filepath.IsAbs(logDir) { + logDir = filepath.Join(workspace, logDir) + } + + // Ensure directory exists with secure permissions + if err := os.MkdirAll(logDir, 0750); err != nil { + return nil, fmt.Errorf("failed to create audit log directory: %w", err) + } + + // Create rotating writer + rw, err := newRotatingWriter(logDir, cfg.Rotation, cfg.Format) + if err != nil { + return nil, fmt.Errorf("failed to create rotating writer: %w", err) + } + + ctx, cancel := context.WithCancel(context.Background()) + + l := &Logger{ + config: cfg, + entries: make(chan *Entry, 1000), // Buffer up to 1000 entries + writer: rw, + ctx: ctx, + cancel: cancel, + filter: newEventFilter(cfg.Events), + } + + // Start background worker + l.wg.Add(1) + go l.worker() + + return l, nil +} + +// Log writes a single audit entry. +// This method is non-blocking; the entry is queued for async writing. +// If the logger is closed or the buffer is full, the entry is dropped. +func (l *Logger) Log(entry *Entry) { + l.mu.RLock() + if l.closed { + l.mu.RUnlock() + return + } + l.mu.RUnlock() + + // Filter by event type + if !l.filter.Allow(entry.EventType) { + return + } + + // Set timestamp if not already set + if entry.Timestamp.IsZero() { + entry.Timestamp = time.Now().UTC() + } + + // Try to queue entry without blocking + select { + case l.entries <- entry: + default: + // Buffer full, drop entry (better than blocking) + // In production, we might want to increment a metric here + } +} + +// LogToolCall logs a tool execution event. +func (l *Logger) LogToolCall(ctx context.Context, data *ToolCallData, duration time.Duration) { + if l == nil { + return + } + + l.Log(&Entry{ + Level: LevelInfo, + Component: "tool", + EventType: EventToolCall, + RequestID: RequestIDFromContext(ctx), + SessionID: SessionIDFromContext(ctx), + AgentID: AgentIDFromContext(ctx), + ToolCall: data, + DurationMs: duration.Milliseconds(), + }) +} + +// LogMessage logs a message event (inbound or outbound). +func (l *Logger) LogMessage(ctx context.Context, direction, contentType, content, messageID string) { + if l == nil { + return + } + + // Truncate content to avoid huge log entries + const maxContentLen = 10000 + if len(content) > maxContentLen { + content = content[:maxContentLen] + "... [truncated]" + } + + l.Log(&Entry{ + Level: LevelInfo, + Component: "channel", + EventType: EventMessage, + RequestID: RequestIDFromContext(ctx), + SessionID: SessionIDFromContext(ctx), + AgentID: AgentIDFromContext(ctx), + Message: &MessageData{ + Direction: direction, + ContentType: contentType, + Content: content, + MessageID: messageID, + }, + }) +} + +// LogError logs an error event. +func (l *Logger) LogError(ctx context.Context, errorType, message string, recoverable bool) { + if l == nil { + return + } + + l.Log(&Entry{ + Level: LevelError, + Component: "system", + EventType: EventError, + RequestID: RequestIDFromContext(ctx), + SessionID: SessionIDFromContext(ctx), + AgentID: AgentIDFromContext(ctx), + Error: &ErrorData{ + ErrorType: errorType, + Message: message, + Recoverable: recoverable, + }, + }) +} + +// LogSystem logs a system-level event. +func (l *Logger) LogSystem(ctx context.Context, operation string, details map[string]interface{}) { + if l == nil { + return + } + + l.Log(&Entry{ + Level: LevelInfo, + Component: "system", + EventType: EventSystem, + RequestID: RequestIDFromContext(ctx), + SessionID: SessionIDFromContext(ctx), + System: &SystemData{ + Operation: operation, + Details: details, + }, + }) +} + +// WithChannelContext returns a context enriched with channel information. +// This context can be used with Log* methods to automatically include channel details. +func WithChannelContext(ctx context.Context, channel, chatID, userID string) context.Context { + ctx = context.WithValue(ctx, channelKey, channel) + ctx = context.WithValue(ctx, chatIDKey, chatID) + ctx = context.WithValue(ctx, userIDKey, userID) + return ctx +} + +// LogWithChannel logs an entry with channel context extracted from the context. +func (l *Logger) LogWithChannel(ctx context.Context, entry *Entry) { + if l == nil { + return + } + + // Extract channel context + if channel, ok := ctx.Value(channelKey).(string); ok { + entry.Channel = channel + } + if chatID, ok := ctx.Value(chatIDKey).(string); ok { + entry.ChatID = chatID + } + if userID, ok := ctx.Value(userIDKey).(string); ok { + entry.UserID = userID + } + + l.Log(entry) +} + +// Close gracefully shuts down the audit logger. +// It flushes any pending entries and closes the log file. +// This method blocks until all pending entries are written. +func (l *Logger) Close() error { + l.mu.Lock() + if l.closed { + l.mu.Unlock() + return nil + } + l.closed = true + l.mu.Unlock() + + // Signal worker to stop + l.cancel() + + // Wait for worker to finish + l.wg.Wait() + + // Close the writer + return l.writer.Close() +} + +// worker processes the entry queue in the background. +func (l *Logger) worker() { + defer l.wg.Done() + + for { + select { + case <-l.ctx.Done(): + // Flush remaining entries + l.flush() + return + case entry := <-l.entries: + if err := l.writeEntry(entry); err != nil { + // Log to stderr as fallback (can't use logger to avoid recursion) + fmt.Fprintf(os.Stderr, "[audit] failed to write entry: %v\n", err) + } + } + } +} + +// flush writes all pending entries without blocking. +func (l *Logger) flush() { + for { + select { + case entry := <-l.entries: + if err := l.writeEntry(entry); err != nil { + fmt.Fprintf(os.Stderr, "[audit] failed to write entry during flush: %v\n", err) + } + default: + // No more entries + return + } + } +} + +// writeEntry serializes and writes a single entry. +func (l *Logger) writeEntry(entry *Entry) error { + var data []byte + var err error + + if l.config.Format == "json" { + data, err = json.Marshal(entry) + if err != nil { + return fmt.Errorf("failed to marshal entry: %w", err) + } + data = append(data, '\n') + } else { + // Text format + data = []byte(l.formatText(entry) + "\n") + } + + return l.writer.Write(data) +} + +// formatText formats an entry as human-readable text. +func (l *Logger) formatText(entry *Entry) string { + return fmt.Sprintf("[%s] %s %s %s: %s", + entry.Timestamp.Format(time.RFC3339), + entry.Level, + entry.Component, + entry.EventType, + l.formatTextDetails(entry), + ) +} + +// formatTextDetails formats event-specific details for text output. +func (l *Logger) formatTextDetails(entry *Entry) string { + switch entry.EventType { + case EventToolCall: + if entry.ToolCall != nil { + return fmt.Sprintf("tool=%s async=%v error=%v", + entry.ToolCall.Name, + entry.ToolCall.IsAsync, + entry.ToolCall.IsError, + ) + } + case EventMessage: + if entry.Message != nil { + return fmt.Sprintf("direction=%s type=%s", + entry.Message.Direction, + entry.Message.ContentType, + ) + } + case EventError: + if entry.Error != nil { + return fmt.Sprintf("type=%s recoverable=%v msg=%s", + entry.Error.ErrorType, + entry.Error.Recoverable, + entry.Error.Message, + ) + } + } + return "" +} + +// channel context keys +type channelContextKey int + +const ( + channelKey channelContextKey = iota + chatIDKey + userIDKey +) + +// eventFilter determines which event types should be logged. +type eventFilter struct { + events config.AuditEvents +} + +func newEventFilter(events config.AuditEvents) *eventFilter { + return &eventFilter{events: events} +} + +func (f *eventFilter) Allow(eventType EventType) bool { + switch eventType { + case EventToolCall: + return f.events.ToolCalls + case EventMessage: + return f.events.Messages + case EventError: + return f.events.Errors + case EventSystem: + return f.events.System + default: + return true + } +} diff --git a/pkg/audit/rotation.go b/pkg/audit/rotation.go new file mode 100644 index 000000000..6ac3d8717 --- /dev/null +++ b/pkg/audit/rotation.go @@ -0,0 +1,276 @@ +package audit + +import ( + "compress/gzip" + "fmt" + "io" + "os" + "path/filepath" + "sort" + "strings" + "sync" + "time" + + "github.com/sipeed/picoclaw/pkg/config" +) + +// rotatingWriter handles file I/O with rotation based on size and age. +// +// Rotation strategy: +// - Daily rotation: New file created at midnight +// - Size-based: File rotated when exceeding MaxSizeMB +// - Cleanup: Old files deleted after MaxAgeDays +// - Backup limit: Only MaxBackups files retained +// - Compression: Old files optionally gzip compressed +// +// The filename pattern is: audit-DDMMYYYY.log[.N][.gz] +type rotatingWriter struct { + baseDir string + baseName string + extension string + rotation config.RotationConfig + format string + + // Current file state + currentFile *os.File + currentSize int64 + currentDate string + + // Synchronization + mu sync.Mutex +} + +// newRotatingWriter creates a new rotating writer. +func newRotatingWriter(baseDir string, rotation config.RotationConfig, format string) (*rotatingWriter, error) { + rw := &rotatingWriter{ + baseDir: baseDir, + baseName: "audit-", + extension: ".log", + rotation: rotation, + format: format, + } + + // Open initial file + if err := rw.rotate(); err != nil { + return nil, err + } + + return rw, nil +} + +// Write writes data to the current log file, rotating if necessary. +func (rw *rotatingWriter) Write(data []byte) error { + rw.mu.Lock() + defer rw.mu.Unlock() + + // Check if rotation needed + today := time.Now().Format("02012006") + needsRotation := false + + if today != rw.currentDate { + // Day changed + needsRotation = true + } else if rw.rotation.MaxSizeMB > 0 { + // Check size limit + maxSize := int64(rw.rotation.MaxSizeMB) * 1024 * 1024 + if rw.currentSize+int64(len(data)) > maxSize { + needsRotation = true + } + } + + if needsRotation { + if err := rw.rotate(); err != nil { + return err + } + } + + // Write data + n, err := rw.currentFile.Write(data) + if err != nil { + return fmt.Errorf("failed to write to log file: %w", err) + } + + rw.currentSize += int64(n) + return nil +} + +// Close closes the current log file. +func (rw *rotatingWriter) Close() error { + rw.mu.Lock() + defer rw.mu.Unlock() + + if rw.currentFile != nil { + return rw.currentFile.Close() + } + return nil +} + +// rotate performs log rotation: +// 1. Close current file +// 2. Compress old file if enabled +// 3. Open new file for today +// 4. Clean up old files +func (rw *rotatingWriter) rotate() error { + // Close current file if open + if rw.currentFile != nil { + if err := rw.currentFile.Close(); err != nil { + return fmt.Errorf("failed to close current log file: %w", err) + } + + // Compress the old file if compression is enabled + if rw.rotation.Compress { + oldPath := rw.currentFile.Name() + go rw.compressFile(oldPath) // Async compression + } + } + + // Update current date + rw.currentDate = time.Now().Format("02012006") + + // Generate new filename + newPath := filepath.Join(rw.baseDir, rw.baseName+rw.currentDate+rw.extension) + + // Check if file already exists (from previous run) + // If so, find next available number + if _, err := os.Stat(newPath); err == nil { + newPath = rw.findNextAvailableName() + } + + // Create new file with secure permissions + file, err := os.OpenFile(newPath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0600) + if err != nil { + return fmt.Errorf("failed to create new log file: %w", err) + } + + rw.currentFile = file + rw.currentSize = 0 + + // Clean up old files asynchronously + go rw.cleanup() + + return nil +} + +// findNextAvailableName finds the next available numbered filename. +func (rw *rotatingWriter) findNextAvailableName() string { + base := filepath.Join(rw.baseDir, rw.baseName+rw.currentDate) + + for i := 1; i < 1000; i++ { + candidate := fmt.Sprintf("%s.%d%s", base, i, rw.extension) + if _, err := os.Stat(candidate); os.IsNotExist(err) { + return candidate + } + } + + // Fallback with timestamp + return fmt.Sprintf("%s.%d%s", base, time.Now().Unix(), rw.extension) +} + +// compressFile compresses a log file with gzip. +func (rw *rotatingWriter) compressFile(path string) { + src, err := os.Open(path) + if err != nil { + return + } + defer src.Close() + + dstPath := path + ".gz" + dst, err := os.Create(dstPath) + if err != nil { + return + } + defer dst.Close() + + gz := gzip.NewWriter(dst) + defer gz.Close() + + if _, err := io.Copy(gz, src); err != nil { + // Clean up partial file + os.Remove(dstPath) + return + } + + // Remove original file after successful compression + src.Close() + os.Remove(path) +} + +// cleanup removes old log files based on rotation settings. +func (rw *rotatingWriter) cleanup() { + files, err := rw.listLogFiles() + if err != nil { + return + } + + // Sort by modification time (oldest first) + sort.Slice(files, func(i, j int) bool { + return files[i].ModTime.Before(files[j].ModTime) + }) + + now := time.Now() + + for _, file := range files { + shouldDelete := false + + // Check age limit + if rw.rotation.MaxAgeDays > 0 { + maxAge := time.Duration(rw.rotation.MaxAgeDays) * 24 * time.Hour + if now.Sub(file.ModTime) > maxAge { + shouldDelete = true + } + } + + // Check backup limit (keep most recent MaxBackups) + if rw.rotation.MaxBackups > 0 && len(files) > rw.rotation.MaxBackups { + shouldDelete = true + } + + if shouldDelete { + os.Remove(file.Path) + // Remove from list so count is accurate + files = files[1:] + } + } +} + +// logFileInfo holds information about a log file for cleanup. +type logFileInfo struct { + Path string + ModTime time.Time +} + +// listLogFiles returns all audit log files in the base directory. +func (rw *rotatingWriter) listLogFiles() ([]logFileInfo, error) { + entries, err := os.ReadDir(rw.baseDir) + if err != nil { + return nil, err + } + + var files []logFileInfo + for _, entry := range entries { + if entry.IsDir() { + continue + } + + name := entry.Name() + // Match audit-*.log or audit-*.log.gz + if !strings.HasPrefix(name, rw.baseName) { + continue + } + if !strings.HasSuffix(name, rw.extension) && !strings.HasSuffix(name, rw.extension+".gz") { + continue + } + + info, err := entry.Info() + if err != nil { + continue + } + + files = append(files, logFileInfo{ + Path: filepath.Join(rw.baseDir, name), + ModTime: info.ModTime(), + }) + } + + return files, nil +} diff --git a/pkg/audit/types.go b/pkg/audit/types.go new file mode 100644 index 000000000..b3b6d668f --- /dev/null +++ b/pkg/audit/types.go @@ -0,0 +1,136 @@ +package audit + +import ( + "context" + "time" +) + +// Entry represents a single audit log entry. +// All fields are JSON-serializable for structured logging. +type Entry struct { + // Core metadata + Timestamp time.Time `json:"timestamp"` + Level Level `json:"level"` + Component string `json:"component"` + EventType EventType `json:"event_type"` + + // Request context for correlation + RequestID string `json:"request_id"` + SessionID string `json:"session_id,omitempty"` + AgentID string `json:"agent_id,omitempty"` + + // Channel context + Channel string `json:"channel,omitempty"` + ChatID string `json:"chat_id,omitempty"` + UserID string `json:"user_id,omitempty"` + + // Event-specific data + ToolCall *ToolCallData `json:"tool_call,omitempty"` + Message *MessageData `json:"message,omitempty"` + Error *ErrorData `json:"error,omitempty"` + System *SystemData `json:"system,omitempty"` + + // Performance metrics + DurationMs int64 `json:"duration_ms,omitempty"` +} + +// Level represents the severity of an audit entry. +type Level string + +const ( + LevelDebug Level = "DEBUG" + LevelInfo Level = "INFO" + LevelWarn Level = "WARN" + LevelError Level = "ERROR" +) + +// EventType categorizes audit entries for filtering and analysis. +type EventType string + +const ( + EventToolCall EventType = "tool_call" + EventMessage EventType = "message" + EventError EventType = "error" + EventSystem EventType = "system" +) + +// ToolCallData captures details of a tool execution. +type ToolCallData struct { + ToolID string `json:"tool_id"` + Name string `json:"name"` + Arguments map[string]interface{} `json:"arguments,omitempty"` + Result string `json:"result,omitempty"` + IsError bool `json:"is_error"` + IsAsync bool `json:"is_async"` +} + +// MessageData captures message flow details. +type MessageData struct { + Direction string `json:"direction"` // "inbound" or "outbound" + ContentType string `json:"content_type"` // "text", "media", "command" + Content string `json:"content,omitempty"` + MessageID string `json:"message_id,omitempty"` +} + +// ErrorData captures error details for debugging. +type ErrorData struct { + ErrorType string `json:"error_type"` + Message string `json:"message"` + StackTrace string `json:"stack_trace,omitempty"` + Recoverable bool `json:"recoverable"` +} + +// SystemData captures system-level events. +type SystemData struct { + Operation string `json:"operation"` + Details map[string]interface{} `json:"details,omitempty"` +} + +// contextKey is a private type for context keys to avoid collisions. +type contextKey int + +const ( + requestIDKey contextKey = iota + sessionIDKey + agentIDKey +) + +// WithRequestID returns a context with the request ID set. +// The request ID is used to correlate all audit entries from a single request. +func WithRequestID(ctx context.Context, requestID string) context.Context { + return context.WithValue(ctx, requestIDKey, requestID) +} + +// RequestIDFromContext extracts the request ID from the context. +func RequestIDFromContext(ctx context.Context) string { + if id, ok := ctx.Value(requestIDKey).(string); ok { + return id + } + return "" +} + +// WithSessionID returns a context with the session ID set. +func WithSessionID(ctx context.Context, sessionID string) context.Context { + return context.WithValue(ctx, sessionIDKey, sessionID) +} + +// SessionIDFromContext extracts the session ID from the context. +func SessionIDFromContext(ctx context.Context) string { + if id, ok := ctx.Value(sessionIDKey).(string); ok { + return id + } + return "" +} + +// WithAgentID returns a context with the agent ID set. +func WithAgentID(ctx context.Context, agentID string) context.Context { + return context.WithValue(ctx, agentIDKey, agentID) +} + +// AgentIDFromContext extracts the agent ID from the context. +func AgentIDFromContext(ctx context.Context) string { + if id, ok := ctx.Value(agentIDKey).(string); ok { + return id + } + return "" +}