Skip to content

Commit b62079d

Browse files
committed
fix: Remove sys___init requirement - sessions always auto-created
1 parent 9349f8d commit b62079d

7 files changed

Lines changed: 40 additions & 46 deletions

File tree

.github/workflows/large-payload-tester.lock.yml

Lines changed: 2 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

.github/workflows/nightly-mcp-stress-test.lock.yml

Lines changed: 2 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

README.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -273,7 +273,7 @@ Available Commands:
273273
Flags:
274274
-c, --config string Path to config file
275275
--config-stdin Read MCP server configuration from stdin (JSON format). When enabled, overrides --config
276-
--enable-difc Enable DIFC enforcement and session requirement (requires sys___init call before tool access)
276+
--enable-difc Enable DIFC enforcement for information flow control
277277
--env string Path to .env file to load environment variables
278278
-h, --help help for awmg
279279
-l, --listen string HTTP server listen address (default "127.0.0.1:3000")

config.example.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -109,7 +109,7 @@ args = [
109109
# ============================================================================
110110

111111
# Enable Data Information Flow Control (DIFC) security model (default: false)
112-
# When true, requires sys___init call before tool access
112+
# When enabled, enforces information flow control policies on tool calls
113113
# This is an experimental feature - keep disabled for standard MCP compatibility
114114
# enable_difc = false
115115

internal/cmd/flags_difc.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ var (
1919

2020
func init() {
2121
RegisterFlag(func(cmd *cobra.Command) {
22-
cmd.Flags().BoolVar(&enableDIFC, "enable-difc", getDefaultEnableDIFC(), "Enable DIFC enforcement and session requirement (requires sys___init call before tool access)")
22+
cmd.Flags().BoolVar(&enableDIFC, "enable-difc", getDefaultEnableDIFC(), "Enable DIFC enforcement for information flow control")
2323
})
2424
}
2525

internal/server/unified.go

Lines changed: 20 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -846,48 +846,35 @@ func (us *UnifiedServer) ensureSessionDirectory(sessionID string) error {
846846
}
847847

848848
// requireSession checks that a session has been initialized for this request
849-
// When DIFC is disabled (default), automatically creates a session if one doesn't exist
849+
// Sessions are automatically created if one doesn't exist (for standard MCP client compatibility)
850850
func (us *UnifiedServer) requireSession(ctx context.Context) error {
851851
sessionID := us.getSessionID(ctx)
852852
log.Printf("Checking session for ID: %s", sessionID)
853853

854-
// If DIFC is disabled (default), use double-checked locking to auto-create session
855-
if !us.enableDIFC {
856-
us.sessionMu.RLock()
857-
session := us.sessions[sessionID]
858-
us.sessionMu.RUnlock()
859-
860-
if session == nil {
861-
// Need to create session - acquire write lock
862-
us.sessionMu.Lock()
863-
// Double-check after acquiring write lock to avoid race condition
864-
if us.sessions[sessionID] == nil {
865-
log.Printf("DIFC disabled: auto-creating session for ID: %s", sessionID)
866-
us.sessions[sessionID] = NewSession(sessionID, "")
867-
log.Printf("Session auto-created for ID: %s", sessionID)
868-
869-
// Ensure session directory exists in payload mount point
870-
// This is done after releasing the lock to avoid holding it during I/O
871-
us.sessionMu.Unlock()
872-
if err := us.ensureSessionDirectory(sessionID); err != nil {
873-
logger.LogWarn("client", "Failed to create session directory for session=%s: %v", sessionID, err)
874-
// Don't fail - payloads will attempt to create the directory when needed
875-
}
876-
return nil
877-
}
878-
us.sessionMu.Unlock()
879-
}
880-
return nil
881-
}
882-
883-
// DIFC is enabled - require explicit session initialization
854+
// Use double-checked locking to auto-create session if needed
884855
us.sessionMu.RLock()
885856
session := us.sessions[sessionID]
886857
us.sessionMu.RUnlock()
887858

888859
if session == nil {
889-
log.Printf("Session not found for ID: %s. Available sessions: %v", sessionID, us.getSessionKeys())
890-
return fmt.Errorf("sys___init must be called before any other tool calls")
860+
// Need to create session - acquire write lock
861+
us.sessionMu.Lock()
862+
// Double-check after acquiring write lock to avoid race condition
863+
if us.sessions[sessionID] == nil {
864+
log.Printf("Auto-creating session for ID: %s", sessionID)
865+
us.sessions[sessionID] = NewSession(sessionID, "")
866+
log.Printf("Session auto-created for ID: %s", sessionID)
867+
868+
// Ensure session directory exists in payload mount point
869+
// This is done after releasing the lock to avoid holding it during I/O
870+
us.sessionMu.Unlock()
871+
if err := us.ensureSessionDirectory(sessionID); err != nil {
872+
logger.LogWarn("client", "Failed to create session directory for session=%s: %v", sessionID, err)
873+
// Don't fail - payloads will attempt to create the directory when needed
874+
}
875+
return nil
876+
}
877+
us.sessionMu.Unlock()
891878
}
892879

893880
log.Printf("Session validated for ID: %s", sessionID)

internal/server/unified_test.go

Lines changed: 13 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -182,10 +182,17 @@ func TestRequireSession(t *testing.T) {
182182
err = us.requireSession(ctxWithSession)
183183
assert.NoError(t, err, "requireSession() failed for valid session")
184184

185-
// Test with invalid session (DIFC enabled)
186-
ctxWithInvalidSession := context.WithValue(ctx, SessionIDContextKey, "invalid-session")
187-
err = us.requireSession(ctxWithInvalidSession)
188-
require.Error(t, err, "requireSession() should fail for invalid session when DIFC is enabled")
185+
// Test with non-existent session - should auto-create even with DIFC enabled
186+
ctxWithNewSession := context.WithValue(ctx, SessionIDContextKey, "new-session")
187+
err = us.requireSession(ctxWithNewSession)
188+
require.NoError(t, err, "requireSession() should auto-create session even when DIFC is enabled")
189+
190+
// Verify session was created
191+
us.sessionMu.RLock()
192+
newSession, exists := us.sessions["new-session"]
193+
us.sessionMu.RUnlock()
194+
require.True(t, exists, "Session should have been auto-created")
195+
require.NotNil(t, newSession, "Session should not be nil")
189196
}
190197

191198
func TestRequireSession_DifcDisabled(t *testing.T) {
@@ -432,8 +439,8 @@ func TestRequireSession_EdgeCases(t *testing.T) {
432439
enableDIFC: true,
433440
sessionID: "nonexistent",
434441
preCreate: false,
435-
wantErr: true,
436-
description: "should deny access to nonexistent session when DIFC enabled",
442+
wantErr: false, // Sessions are auto-created regardless of DIFC setting
443+
description: "should auto-create session even when DIFC enabled",
437444
},
438445
{
439446
name: "DIFC disabled without session",

0 commit comments

Comments
 (0)