From fd32e4a97f3eebacc99e9caac2191340f277feb4 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 14 Apr 2026 18:15:41 +0000 Subject: [PATCH 1/6] Initial plan From 98d5a84908568e85cce745726c334ff8e215d68f Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 14 Apr 2026 18:25:29 +0000 Subject: [PATCH 2/6] wazero improvements: Registry.Close, logging namespace fix, typed trap detection - Add Registry.Close(ctx) method to close all registered guards that implement io.Closer; call it from UnifiedServer.Close() and InitiateShutdown() to release WASM runtime resources on shutdown - Fix logging namespace confusion in registry.go: replace log.Printf (which was using the guard:context namespace logger) with logger.LogInfo("guard", ...) for operational events - Use typed sys.ExitError check in isWasmTrap: check exit code 0 as a normal process exit (not a trap) before falling back to string matching - Add tests for Registry.Close and sys.ExitError handling in isWasmTrap Agent-Logs-Url: https://github.com/github/gh-aw-mcpg/sessions/842dba68-e351-47f7-84e1-cf101f122182 Co-authored-by: lpcox <15877973+lpcox@users.noreply.github.com> --- internal/guard/guard_test.go | 66 ++++++++++++++++++++++++++++++++++++ internal/guard/registry.go | 22 +++++++++--- internal/guard/wasm.go | 19 +++++++++-- internal/guard/wasm_test.go | 21 ++++++++++++ internal/server/unified.go | 8 +++++ 5 files changed, 129 insertions(+), 7 deletions(-) diff --git a/internal/guard/guard_test.go b/internal/guard/guard_test.go index 93b85d9f..07e05e98 100644 --- a/internal/guard/guard_test.go +++ b/internal/guard/guard_test.go @@ -2,6 +2,7 @@ package guard import ( "context" + "errors" "sync" "testing" @@ -17,6 +18,18 @@ type mockGuard struct { id string } +// mockClosableGuard is a guard that tracks whether Close was called +type mockClosableGuard struct { + mockGuard + closed bool + closeErr error +} + +func (m *mockClosableGuard) Close(ctx context.Context) error { + m.closed = true + return m.closeErr +} + func (m *mockGuard) Name() string { return "mock-" + m.id } func (m *mockGuard) LabelAgent(ctx context.Context, policy interface{}, backend BackendCaller, caps *difc.Capabilities) (*LabelAgentResult, error) { return &LabelAgentResult{DIFCMode: difc.ModeStrict}, nil @@ -450,6 +463,59 @@ func TestGuardRegistry_HasNonNoopGuard(t *testing.T) { }) } +func TestGuardRegistry_Close(t *testing.T) { + t.Run("close calls Close on guards that implement it", func(t *testing.T) { + registry := NewRegistry() + g := &mockClosableGuard{mockGuard: mockGuard{id: "wasm"}} + registry.Register("server1", g) + + registry.Close(context.Background()) + + assert.True(t, g.closed, "expected guard Close to be called") + }) + + t.Run("close skips guards that do not implement Close", func(t *testing.T) { + registry := NewRegistry() + registry.Register("server1", NewNoopGuard()) + + // Should not panic + registry.Close(context.Background()) + }) + + t.Run("close on empty registry is safe", func(t *testing.T) { + registry := NewRegistry() + // Should not panic + registry.Close(context.Background()) + }) + + t.Run("close calls Close on all closable guards", func(t *testing.T) { + registry := NewRegistry() + g1 := &mockClosableGuard{mockGuard: mockGuard{id: "wasm1"}} + g2 := &mockClosableGuard{mockGuard: mockGuard{id: "wasm2"}} + registry.Register("server1", g1) + registry.Register("server2", g2) + + registry.Close(context.Background()) + + assert.True(t, g1.closed, "expected guard 1 Close to be called") + assert.True(t, g2.closed, "expected guard 2 Close to be called") + }) + + t.Run("close continues when one guard returns an error", func(t *testing.T) { + registry := NewRegistry() + g1 := &mockClosableGuard{mockGuard: mockGuard{id: "failing"}, closeErr: errors.New("close failed")} + g2 := &mockClosableGuard{mockGuard: mockGuard{id: "ok"}} + registry.Register("server1", g1) + registry.Register("server2", g2) + + // Should not panic even when one guard returns an error + registry.Close(context.Background()) + + assert.True(t, g1.closed, "expected failing guard Close to be called") + assert.True(t, g2.closed, "expected ok guard Close to be called") + }) +} + func TestCreateGuard(t *testing.T) { tests := []struct { name string diff --git a/internal/guard/registry.go b/internal/guard/registry.go index 08143455..4855d8b5 100644 --- a/internal/guard/registry.go +++ b/internal/guard/registry.go @@ -1,6 +1,7 @@ package guard import ( + "context" "fmt" "sync" @@ -30,7 +31,7 @@ func (r *Registry) Register(serverID string, guard Guard) { defer r.mu.Unlock() r.guards[serverID] = guard - log.Printf("[Guard] Registered guard '%s' for server '%s'", guard.Name(), serverID) + logger.LogInfo("guard", "Registered guard '%s' for server '%s'", guard.Name(), serverID) } // Get retrieves the guard for a server, or returns a noop guard if not found @@ -46,7 +47,6 @@ func (r *Registry) Get(serverID string) Guard { // Return noop guard as default debugLog.Printf("No guard registered for serverID=%s, returning noop guard", serverID) - log.Printf("[Guard] No guard registered for server '%s', using noop guard", serverID) return NewNoopGuard() } @@ -76,7 +76,7 @@ func (r *Registry) Remove(serverID string) { r.mu.Lock() defer r.mu.Unlock() delete(r.guards, serverID) - log.Printf("[Guard] Removed guard for server '%s'", serverID) + logger.LogInfo("guard", "Removed guard for server '%s'", serverID) } // List returns all registered server IDs @@ -103,6 +103,20 @@ func (r *Registry) GetGuardInfo() map[string]string { return info } +// Close closes all registered guards that implement io.Closer. +// It should be called during server shutdown to release WASM runtime resources. +func (r *Registry) Close(ctx context.Context) { + r.mu.Lock() + defer r.mu.Unlock() + for id, g := range r.guards { + if c, ok := g.(interface{ Close(context.Context) error }); ok { + if err := c.Close(ctx); err != nil { + logger.LogWarn("guard", "Failed to close guard for server %s: %v", id, err) + } + } + } +} + // GuardFactory is a function that creates a guard instance type GuardFactory func() (Guard, error) @@ -116,7 +130,7 @@ func RegisterGuardType(name string, factory GuardFactory) { registeredGuardsMu.Lock() defer registeredGuardsMu.Unlock() registeredGuards[name] = factory - log.Printf("[Guard] Registered guard type: %s", name) + logger.LogInfo("guard", "Registered guard type: %s", name) } // CreateGuard creates a guard instance by name using registered factories diff --git a/internal/guard/wasm.go b/internal/guard/wasm.go index ea144fa9..2a300d59 100644 --- a/internal/guard/wasm.go +++ b/internal/guard/wasm.go @@ -15,6 +15,7 @@ import ( "github.com/tetratelabs/wazero" "github.com/tetratelabs/wazero/api" "github.com/tetratelabs/wazero/imports/wasi_snapshot_preview1" + "github.com/tetratelabs/wazero/sys" ) var logWasm = logger.New("guard:wasm") @@ -830,10 +831,22 @@ func parsePathLabeledResponse(responseJSON []byte, originalData interface{}) (di return pld.ToCollectionLabeledData(), nil } -// isWasmTrap reports whether err is a WASM execution trap such as the -// "wasm error: unreachable" produced when a Rust-compiled guard panics. +// isWasmTrap reports whether err represents a WASM execution trap that should +// permanently poison the guard. Normal process exits (exit code 0, e.g. TinyGo +// init) are NOT considered traps. A non-zero exit code is treated as a trap. +// As a fallback for wazero execution faults (e.g. Rust panic → unreachable), +// the function also matches on wazero's "wasm error:" message prefix. func isWasmTrap(err error) bool { - return err != nil && strings.Contains(err.Error(), "wasm error:") + if err == nil { + return false + } + // A normal WASI process exit (exit code 0) is not a trap — don't poison the guard. + var exitErr *sys.ExitError + if errors.As(err, &exitErr) { + return exitErr.ExitCode() != 0 + } + // Fallback for wazero execution traps (e.g. Rust panic → unreachable). + return strings.Contains(err.Error(), "wasm error:") } // callWasmFunction calls an exported function in the WASM module. diff --git a/internal/guard/wasm_test.go b/internal/guard/wasm_test.go index 66314b2d..776501fe 100644 --- a/internal/guard/wasm_test.go +++ b/internal/guard/wasm_test.go @@ -16,6 +16,7 @@ import ( "github.com/github/gh-aw-mcpg/internal/difc" "github.com/tetratelabs/wazero" + "github.com/tetratelabs/wazero/sys" ) func TestMain(m *testing.M) { @@ -1152,6 +1153,26 @@ func TestIsWasmTrap(t *testing.T) { err := errors.New("wasm error: out of bounds memory access") assert.True(t, isWasmTrap(err)) }) + + t.Run("sys.ExitError with exit code 0 is not a trap", func(t *testing.T) { + err := sys.NewExitError(0) + assert.False(t, isWasmTrap(err)) + }) + + t.Run("sys.ExitError with non-zero exit code is a trap", func(t *testing.T) { + err := sys.NewExitError(1) + assert.True(t, isWasmTrap(err)) + }) + + t.Run("wrapped sys.ExitError with exit code 0 is not a trap", func(t *testing.T) { + err := fmt.Errorf("wrapper: %w", sys.NewExitError(0)) + assert.False(t, isWasmTrap(err)) + }) + + t.Run("wrapped sys.ExitError with non-zero exit code is a trap", func(t *testing.T) { + err := fmt.Errorf("wrapper: %w", sys.NewExitError(2)) + assert.True(t, isWasmTrap(err)) + }) } func TestWasmGuardFailedState(t *testing.T) { diff --git a/internal/server/unified.go b/internal/server/unified.go index 7c0f5ce8..b622b849 100644 --- a/internal/server/unified.go +++ b/internal/server/unified.go @@ -718,6 +718,9 @@ func (us *UnifiedServer) GetToolHandler(backendID string, toolName string) func( // Close cleans up resources func (us *UnifiedServer) Close() error { + if us.guardRegistry != nil { + us.guardRegistry.Close(context.Background()) + } us.launcher.Close() return nil } @@ -753,6 +756,11 @@ func (us *UnifiedServer) InitiateShutdown() int { logger.LogInfo("shutdown", "Terminating %d backend servers", serversTerminated) us.launcher.Close() + // Release WASM runtime resources held by guards + if us.guardRegistry != nil { + us.guardRegistry.Close(context.Background()) + } + logger.LogInfo("shutdown", "Backend servers terminated successfully") }) return serversTerminated From 995beeb80966b351c47b0123f885338f94d149d0 Mon Sep 17 00:00:00 2001 From: Landon Cox Date: Tue, 14 Apr 2026 11:41:01 -0700 Subject: [PATCH 3/6] Update internal/guard/registry.go Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- internal/guard/registry.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/internal/guard/registry.go b/internal/guard/registry.go index 4855d8b5..d5849087 100644 --- a/internal/guard/registry.go +++ b/internal/guard/registry.go @@ -103,7 +103,7 @@ func (r *Registry) GetGuardInfo() map[string]string { return info } -// Close closes all registered guards that implement io.Closer. +// Close closes all registered guards that implement Close(context.Context) error. // It should be called during server shutdown to release WASM runtime resources. func (r *Registry) Close(ctx context.Context) { r.mu.Lock() From 6b43c7a5492b5a60d9c3968d598b23c358eb23ab Mon Sep 17 00:00:00 2001 From: Landon Cox Date: Tue, 14 Apr 2026 11:41:14 -0700 Subject: [PATCH 4/6] Update internal/guard/registry.go Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- internal/guard/registry.go | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/internal/guard/registry.go b/internal/guard/registry.go index d5849087..e1033615 100644 --- a/internal/guard/registry.go +++ b/internal/guard/registry.go @@ -106,13 +106,23 @@ func (r *Registry) GetGuardInfo() map[string]string { // Close closes all registered guards that implement Close(context.Context) error. // It should be called during server shutdown to release WASM runtime resources. func (r *Registry) Close(ctx context.Context) { + type closableGuard struct { + id string + c interface{ Close(context.Context) error } + } + r.mu.Lock() - defer r.mu.Unlock() + closers := make([]closableGuard, 0, len(r.guards)) for id, g := range r.guards { if c, ok := g.(interface{ Close(context.Context) error }); ok { - if err := c.Close(ctx); err != nil { - logger.LogWarn("guard", "Failed to close guard for server %s: %v", id, err) - } + closers = append(closers, closableGuard{id: id, c: c}) + } + } + r.mu.Unlock() + + for _, guard := range closers { + if err := guard.c.Close(ctx); err != nil { + logger.LogWarn("guard", "Failed to close guard for server %s: %v", guard.id, err) } } } From 64bf699f75dcc05b4320ec982344ff2e09d278f6 Mon Sep 17 00:00:00 2001 From: Landon Cox Date: Tue, 14 Apr 2026 11:41:25 -0700 Subject: [PATCH 5/6] Update internal/server/unified.go Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- internal/server/unified.go | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/internal/server/unified.go b/internal/server/unified.go index b622b849..ce433ecb 100644 --- a/internal/server/unified.go +++ b/internal/server/unified.go @@ -718,10 +718,7 @@ func (us *UnifiedServer) GetToolHandler(backendID string, toolName string) func( // Close cleans up resources func (us *UnifiedServer) Close() error { - if us.guardRegistry != nil { - us.guardRegistry.Close(context.Background()) - } - us.launcher.Close() + us.InitiateShutdown() return nil } From 13ddd06d55a2fde21cd21e4cf43dae7da0b04e91 Mon Sep 17 00:00:00 2001 From: Landon Cox Date: Tue, 14 Apr 2026 11:51:44 -0700 Subject: [PATCH 6/6] review: use RLock in Registry.Close, add shutdown log, test double-close - Use RLock/RUnlock instead of Lock/Unlock in Registry.Close() since the guards map is only read (not modified) during the closers scan - Add summary LogInfo after closing guards for shutdown visibility - Add double-close idempotency test with closeCount tracking on mock Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- internal/guard/guard_test.go | 17 +++++++++++++++-- internal/guard/registry.go | 7 +++++-- 2 files changed, 20 insertions(+), 4 deletions(-) diff --git a/internal/guard/guard_test.go b/internal/guard/guard_test.go index 07e05e98..47a8426a 100644 --- a/internal/guard/guard_test.go +++ b/internal/guard/guard_test.go @@ -21,12 +21,14 @@ type mockGuard struct { // mockClosableGuard is a guard that tracks whether Close was called type mockClosableGuard struct { mockGuard - closed bool - closeErr error + closed bool + closeCount int + closeErr error } func (m *mockClosableGuard) Close(ctx context.Context) error { m.closed = true + m.closeCount++ return m.closeErr } @@ -514,6 +516,17 @@ func TestGuardRegistry_Close(t *testing.T) { assert.True(t, g1.closed, "expected failing guard Close to be called") assert.True(t, g2.closed, "expected ok guard Close to be called") }) + + t.Run("double close is safe", func(t *testing.T) { + registry := NewRegistry() + g := &mockClosableGuard{mockGuard: mockGuard{id: "wasm"}} + registry.Register("server1", g) + + registry.Close(context.Background()) + registry.Close(context.Background()) + + assert.Equal(t, 2, g.closeCount, "Close should be called twice without panic") + }) } func TestCreateGuard(t *testing.T) { diff --git a/internal/guard/registry.go b/internal/guard/registry.go index e1033615..44bfceb4 100644 --- a/internal/guard/registry.go +++ b/internal/guard/registry.go @@ -111,20 +111,23 @@ func (r *Registry) Close(ctx context.Context) { c interface{ Close(context.Context) error } } - r.mu.Lock() + r.mu.RLock() closers := make([]closableGuard, 0, len(r.guards)) for id, g := range r.guards { if c, ok := g.(interface{ Close(context.Context) error }); ok { closers = append(closers, closableGuard{id: id, c: c}) } } - r.mu.Unlock() + r.mu.RUnlock() for _, guard := range closers { if err := guard.c.Close(ctx); err != nil { logger.LogWarn("guard", "Failed to close guard for server %s: %v", guard.id, err) } } + if len(closers) > 0 { + logger.LogInfo("guard", "Closed %d guard(s)", len(closers)) + } } // GuardFactory is a function that creates a guard instance