Close workflow audit log writers

pkg/audit/auditor.go

@@ -107,7 +107,15 @@ func NewAuditorWithTransport(config *Config, transportType string) (*Auditor, er
 // Close closes the underlying log writer if it implements io.Closer.
 // This should be called when the auditor is no longer needed to properly release resources.
 func (a *Auditor) Close() error {
-	if closer, ok := a.logWriter.(io.Closer); ok {
+	return closeLogWriter(a.logWriter)
+}
+
+func closeLogWriter(logWriter io.Writer) error {
+	if logWriter == os.Stdout || logWriter == os.Stderr {
+		return nil
+	}
+
+	if closer, ok := logWriter.(io.Closer); ok {
 		return closer.Close()
 	}
 	return nil

pkg/audit/auditor_test.go

@@ -11,6 +11,7 @@ import (
 	"io"
 	"net/http"
 	"net/http/httptest"
+	"os"
 	"strings"
 	"testing"
 	"time"
@@ -33,6 +34,16 @@ func TestNewAuditor(t *testing.T) {
 	assert.Equal(t, config, auditor.config)
 }
 
+func TestAuditor_CloseDoesNotCloseStdout(t *testing.T) {
+	t.Parallel()
+
+	auditor := &Auditor{logWriter: os.Stdout}
+
+	require.NoError(t, auditor.Close())
+	_, err := os.Stdout.Write(nil)
+	require.NoError(t, err, "Close() must not close os.Stdout")
+}
+
 func TestAuditorMiddlewareDisabled(t *testing.T) {
 	t.Parallel()
 	config := &Config{}

pkg/audit/workflow_auditor.go

@@ -8,6 +8,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"io"
 	"log/slog"
 	"time"
 
@@ -21,6 +22,7 @@ type WorkflowAuditor struct {
 	auditLogger *slog.Logger
 	config      *Config
 	component   string
+	logWriter   io.Writer
 }
 
 // NewWorkflowAuditor creates a new workflow auditor.
@@ -45,9 +47,16 @@ func NewWorkflowAuditor(config *Config) (*WorkflowAuditor, error) {
 		auditLogger: NewAuditLogger(logWriter),
 		config:      config,
 		component:   component,
+		logWriter:   logWriter,
 	}, nil
 }
 
+// Close closes the underlying log writer if it owns a closeable resource.
+// This should be called when the workflow auditor is no longer needed.
+func (w *WorkflowAuditor) Close() error {
+	return closeLogWriter(w.logWriter)
+}
+
 // LogWorkflowStarted logs the start of workflow execution.
 func (w *WorkflowAuditor) LogWorkflowStarted(
 	ctx context.Context,

pkg/audit/workflow_auditor_test.go

@@ -23,11 +23,25 @@ type testLogWriter struct {
 	logs []string
 }
 
+type closeTrackingWriter struct {
+	closed   bool
+	closeErr error
+}
+
 func (w *testLogWriter) Write(p []byte) (n int, err error) {
 	w.logs = append(w.logs, string(p))
 	return len(p), nil
 }
 
+func (*closeTrackingWriter) Write(p []byte) (n int, err error) {
+	return len(p), nil
+}
+
+func (w *closeTrackingWriter) Close() error {
+	w.closed = true
+	return w.closeErr
+}
+
 func (w *testLogWriter) getLastLog() string {
 	if len(w.logs) == 0 {
 		return ""
@@ -52,6 +66,7 @@ func createTestAuditor(t *testing.T, config *Config) (*WorkflowAuditor, *testLog
 		auditLogger: NewAuditLogger(writer),
 		config:      config,
 		component:   "vmcp-composer",
+		logWriter:   writer,
 	}
 
 	return auditor, writer
@@ -122,6 +137,48 @@ func TestNewWorkflowAuditor(t *testing.T) {
 	}
 }
 
+func TestWorkflowAuditor_Close(t *testing.T) {
+	t.Parallel()
+
+	t.Run("closes retained file writer", func(t *testing.T) {
+		t.Parallel()
+
+		logFilePath := t.TempDir() + "/workflow-audit.log"
+		auditor, err := NewWorkflowAuditor(&Config{LogFile: logFilePath})
+		require.NoError(t, err)
+
+		_, ok := auditor.logWriter.(interface{ Close() error })
+		require.True(t, ok, "file-backed workflow auditor should retain a closeable writer")
+
+		require.NoError(t, auditor.Close())
+	})
+
+	t.Run("does not close stdout", func(t *testing.T) {
+		t.Parallel()
+
+		auditor, err := NewWorkflowAuditor(&Config{})
+		require.NoError(t, err)
+
+		require.NoError(t, auditor.Close())
+		_, err = os.Stdout.Write(nil)
+		require.NoError(t, err, "Close() must not close os.Stdout")
+		assert.Same(t, os.Stdout, auditor.logWriter)
+	})
+
+	t.Run("propagates close errors", func(t *testing.T) {
+		t.Parallel()
+
+		closeErr := errors.New("close failed")
+		writer := &closeTrackingWriter{closeErr: closeErr}
+		auditor := &WorkflowAuditor{logWriter: writer}
+
+		err := auditor.Close()
+
+		require.ErrorIs(t, err, closeErr)
+		assert.True(t, writer.closed)
+	})
+}
+
 func TestWorkflowAuditor_LogWorkflowStarted(t *testing.T) {
 	t.Parallel()
 

pkg/vmcp/core/core_vmcp.go

@@ -71,6 +71,10 @@ type coreVMCP struct {
 	// by advertised tool name.
 	workflowDefs map[string]*composer.WorkflowDefinition
 
+	// workflowAuditor owns the optional workflow audit log writer and is closed
+	// with the core. Nil when workflow audit logging is disabled.
+	workflowAuditor *audit.WorkflowAuditor
+
 	// composerFactory builds a per-call composite-tool engine bound to a routing
 	// table, generalizing server.New's sessionComposerFactory (server.go:393).
 	composerFactory func(sessionRT *vmcp.RoutingTable, sessionTools []vmcp.Tool) composer.Composer
@@ -138,6 +142,14 @@ func New(cfg *Config) (VMCP, error) {
 		}
 		slog.Info("workflow audit logging enabled")
 	}
+	closeWorkflowAuditor := func() {
+		if workflowAuditor == nil {
+			return
+		}
+		if err := workflowAuditor.Close(); err != nil {
+			slog.Warn("failed to close workflow auditor", "error", err)
+		}
+	}
 
 	// The elicitation handler depends only on the domain-typed ElicitationRequester
 	// (#5436); no mcp-go types cross this boundary (vmcp anti-pattern #5).
@@ -168,6 +180,7 @@ func New(cfg *Config) (VMCP, error) {
 	instruments, err := newWorkflowInstruments(cfg.TelemetryProvider)
 	if err != nil {
 		stopStore()
+		closeWorkflowAuditor()
 		return nil, fmt.Errorf("failed to create workflow telemetry instruments: %w", err)
 	}
 
@@ -195,6 +208,7 @@ func New(cfg *Config) (VMCP, error) {
 	workflowDefs, err := validateWorkflowDefs(validationEngine, cfg.WorkflowDefs)
 	if err != nil {
 		stopStore()
+		closeWorkflowAuditor()
 		return nil, fmt.Errorf("workflow validation failed: %w", err)
 	}
 
@@ -206,6 +220,7 @@ func New(cfg *Config) (VMCP, error) {
 	healthMonitor, healthProvider, err := buildHealthMonitor(cfg)
 	if err != nil {
 		stopStore()
+		closeWorkflowAuditor()
 		return nil, err
 	}
 
@@ -217,6 +232,7 @@ func New(cfg *Config) (VMCP, error) {
 		healthMonitor:   healthMonitor,
 		admission:       admission,
 		workflowDefs:    workflowDefs,
+		workflowAuditor: workflowAuditor,
 		composerFactory: composerFactory,
 		stopStore:       stopStore,
 	}, nil
@@ -543,6 +559,11 @@ func (c *coreVMCP) InvalidateCapabilityCache() {
 func (c *coreVMCP) Close() error {
 	c.closeOnce.Do(func() {
 		c.stopStore()
+		if c.workflowAuditor != nil {
+			if err := c.workflowAuditor.Close(); err != nil {
+				slog.Warn("failed to close workflow auditor", "error", err)
+			}
+		}
 		if c.healthMonitor != nil {
 			if err := c.healthMonitor.Stop(); err != nil {
 				slog.Warn("failed to stop health monitor", "error", err)

pkg/vmcp/core/core_vmcp_test.go

@@ -8,13 +8,16 @@ import (
 	"context"
 	"errors"
 	"log/slog"
+	"os"
+	"path/filepath"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/mock/gomock"
 
+	"github.com/stacklok/toolhive/pkg/audit"
 	"github.com/stacklok/toolhive/pkg/auth"
 	"github.com/stacklok/toolhive/pkg/vmcp"
 	"github.com/stacklok/toolhive/pkg/vmcp/aggregator"
@@ -61,6 +64,27 @@ func baseConfig(t *testing.T) (*Config, *coreMocks) {
 	return cfg, m
 }
 
+func requireNoOpenFDForPath(t *testing.T, path string) {
+	t.Helper()
+
+	entries, err := os.ReadDir("/proc/self/fd")
+	require.NoError(t, err)
+	for _, entry := range entries {
+		target, err := os.Readlink(filepath.Join("/proc/self/fd", entry.Name()))
+		if err != nil {
+			continue
+		}
+		require.NotEqual(t, path, target, "audit log file descriptor must be closed")
+	}
+}
+
+func closeErrorPathAuditConfig(t *testing.T) (*audit.Config, string) {
+	t.Helper()
+
+	path := filepath.Join(t.TempDir(), "workflow-audit.log")
+	return &audit.Config{LogFile: path}, path
+}
+
 // testBackendID is the single backend ID used across these tests.
 const testBackendID = "be1"
 
@@ -115,6 +139,84 @@ func TestNew_NilConfig(t *testing.T) {
 	assert.Nil(t, c)
 }
 
+func TestNew_CloseReleasesWorkflowAuditLogFile(t *testing.T) {
+	t.Parallel()
+
+	t.Run("closes retained audit log file", func(t *testing.T) {
+		t.Parallel()
+
+		cfg, _ := baseConfig(t)
+		cfg.AuditConfig, _ = closeErrorPathAuditConfig(t)
+
+		c, err := New(cfg)
+		require.NoError(t, err)
+		require.NoError(t, c.Close())
+
+		require.ErrorIs(t, c.(*coreVMCP).workflowAuditor.Close(), os.ErrClosed)
+	})
+
+	t.Run("logs close errors without failing Close", func(t *testing.T) {
+		t.Parallel()
+
+		cfg, _ := baseConfig(t)
+		cfg.AuditConfig, _ = closeErrorPathAuditConfig(t)
+
+		c, err := New(cfg)
+		require.NoError(t, err)
+		core := c.(*coreVMCP)
+		require.NoError(t, core.workflowAuditor.Close())
+
+		require.NoError(t, c.Close())
+	})
+}
+
+func TestNew_ErrorPathsCloseWorkflowAuditLogFile(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name   string
+		mutate func(*testing.T, *Config, *coreMocks)
+	}{
+		{
+			name: "workflow validation error",
+			mutate: func(_ *testing.T, cfg *Config, _ *coreMocks) {
+				cfg.WorkflowDefs = map[string]*composer.WorkflowDefinition{
+					"wf": {
+						Name: "wf",
+						Steps: []composer.WorkflowStep{
+							{ID: "s1", Type: composer.StepTypeTool, Tool: "be1.tool", DependsOn: []string{"s2"}},
+							{ID: "s2", Type: composer.StepTypeTool, Tool: "be1.tool", DependsOn: []string{"s1"}},
+						},
+					},
+				}
+			},
+		},
+		{
+			name: "health monitor creation error",
+			mutate: func(_ *testing.T, cfg *Config, mocks *coreMocks) {
+				mocks.reg.EXPECT().List(gomock.Any()).Return(nil)
+				cfg.HealthMonitorConfig = &health.MonitorConfig{}
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+			cfg, mocks := baseConfig(t)
+			var auditLogPath string
+			cfg.AuditConfig, auditLogPath = closeErrorPathAuditConfig(t)
+			tt.mutate(t, cfg, mocks)
+
+			c, err := New(cfg)
+
+			require.Error(t, err)
+			assert.Nil(t, c)
+			requireNoOpenFDForPath(t, auditLogPath)
+		})
+	}
+}
+
 func TestNew_ValidatesWorkflows(t *testing.T) {
 	t.Parallel()