mcp: add sequential thinking server example (#51)

Add example of implementing mcp on [sequentialthinking
server](https://github.com/modelcontextprotocol/servers/tree/main/src/sequentialthinking),
from [Example Servers](https://modelcontextprotocol.io/examples).

The PR includes dynamic and reflective problem-solving through
structured thinking processes.
The server provides tools for:
- starting thinking sessions
- adding sequential thoughts with step tracking
- revising previous thoughts
- creating alternative reasoning branches
-  reviewing complete processes.

Also, features include thread-safe session management, adaptive planning
that adjusts step counts dynamically.

examples/sequentialthinking/README.md

@@ -0,0 +1,203 @@
+# Sequential Thinking MCP Server
+
+This example shows a Model Context Protocol (MCP) server that enables dynamic and reflective problem-solving through structured thinking processes. It helps break down complex problems into manageable, sequential thought steps with support for revision and branching.
+
+## Features
+
+The server provides three main tools for managing thinking sessions:
+
+### 1. Start Thinking (`start_thinking`)
+
+Begins a new sequential thinking session for a complex problem.
+
+**Parameters:**
+
+- `problem` (string): The problem or question to think about
+- `sessionId` (string, optional): Custom session identifier
+- `estimatedSteps` (int, optional): Initial estimate of thinking steps needed
+
+### 2. Continue Thinking (`continue_thinking`)
+
+Adds the next thought step, revises previous steps, or creates alternative branches.
+
+**Parameters:**
+
+- `sessionId` (string): The thinking session to continue
+- `thought` (string): The current thought or analysis
+- `nextNeeded` (bool, optional): Whether another thinking step is needed
+- `reviseStep` (int, optional): Step number to revise (1-based)
+- `createBranch` (bool, optional): Create an alternative reasoning path
+- `estimatedTotal` (int, optional): Update total estimated steps
+
+### 3. Review Thinking (`review_thinking`)
+
+Provides a complete review of the thinking process for a session.
+
+**Parameters:**
+
+- `sessionId` (string): The session to review
+
+## Resources
+
+### Thinking History (`thinking://sessions` or `thinking://{sessionId}`)
+
+Access thinking session data and history in JSON format.
+
+- `thinking://sessions` - List all thinking sessions
+- `thinking://{sessionId}` - Get specific session details
+
+## Core Concepts
+
+### Sequential Processing
+
+Problems are broken down into numbered thought steps that build upon each other, maintaining context and allowing for systematic analysis.
+
+### Dynamic Revision
+
+Any previous thought step can be revised and updated, with the system tracking which thoughts have been modified.
+
+### Alternative Branching
+
+Create alternative reasoning paths to explore different approaches to the same problem, allowing for comparative analysis.
+
+### Adaptive Planning
+
+The estimated number of thinking steps can be adjusted dynamically as understanding of the problem evolves.
+
+## Running the Server
+
+### Standard I/O Mode
+
+```bash
+go run .
+```
+
+### HTTP Mode  
+
+```bash
+go run . -http :8080
+```
+
+## Example Usage
+
+### Starting a Thinking Session
+
+```json
+{
+  "method": "tools/call",
+  "params": {
+    "name": "start_thinking",
+    "arguments": {
+      "problem": "How should I design a scalable microservices architecture?",
+      "sessionId": "architecture_design",
+      "estimatedSteps": 8
+    }
+  }
+}
+```
+
+### Adding Sequential Thoughts
+
+```json
+{
+  "method": "tools/call", 
+  "params": {
+    "name": "continue_thinking",
+    "arguments": {
+      "sessionId": "architecture_design",
+      "thought": "First, I need to identify the core business domains and their boundaries to determine service decomposition."
+    }
+  }
+}
+```
+
+### Revising a Previous Step
+
+```json
+{
+  "method": "tools/call",
+  "params": {
+    "name": "continue_thinking", 
+    "arguments": {
+      "sessionId": "architecture_design",
+      "thought": "Actually, before identifying domains, I should analyze the current system's pain points and requirements.",
+      "reviseStep": 1
+    }
+  }
+}
+```
+
+### Creating an Alternative Branch
+
+```json
+{
+  "method": "tools/call",
+  "params": {
+    "name": "continue_thinking",
+    "arguments": {
+      "sessionId": "architecture_design", 
+      "thought": "Alternative approach: Start with a monolith-first strategy and extract services gradually.",
+      "createBranch": true
+    }
+  }
+}
+```
+
+### Completing the Thinking Process
+
+```json
+{
+  "method": "tools/call",
+  "params": {
+    "name": "continue_thinking",
+    "arguments": {
+      "sessionId": "architecture_design",
+      "thought": "Based on this analysis, I recommend starting with 3 core services: User Management, Order Processing, and Inventory Management.",
+      "nextNeeded": false
+    }
+  }
+}
+```
+
+### Reviewing the Complete Process
+
+```json
+{
+  "method": "tools/call",
+  "params": {
+    "name": "review_thinking", 
+    "arguments": {
+      "sessionId": "architecture_design"
+    }
+  }
+}
+```
+
+## Session State Management
+
+Each thinking session maintains:
+
+- **Session metadata**: ID, problem statement, creation time, current status
+- **Thought sequence**: Ordered list of thoughts with timestamps and revision history  
+- **Progress tracking**: Current step and estimated total steps
+- **Branch relationships**: Links to alternative reasoning paths
+- **Status management**: Active, completed, or paused sessions
+
+## Use Cases
+
+**Ideal for:**
+
+- Complex problem analysis requiring step-by-step breakdown
+- Design decisions needing systematic evaluation
+- Scenarios where initial scope is unclear and may evolve
+- Problems requiring alternative approach exploration
+- Situations needing detailed reasoning documentation
+
+**Examples:**
+
+- Software architecture design
+- Research methodology planning  
+- Strategic business decisions
+- Technical troubleshooting
+- Creative problem solving
+- Academic research planning

examples/sequentialthinking/main.go

@@ -0,0 +1,544 @@
+// Copyright 2025 The Go MCP SDK Authors. All rights reserved.
+// Use of this source code is governed by an MIT-style
+// license that can be found in the LICENSE file.
+
+package main
+
+import (
+	"context"
+	"crypto/rand"
+	"encoding/json"
+	"flag"
+	"fmt"
+	"log"
+	"maps"
+	"net/http"
+	"net/url"
+	"os"
+	"slices"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/modelcontextprotocol/go-sdk/mcp"
+)
+
+var httpAddr = flag.String("http", "", "if set, use streamable HTTP at this address, instead of stdin/stdout")
+
+// A Thought is a single step in the thinking process.
+type Thought struct {
+	// Index of the thought within the session (1-based).
+	Index int `json:"index"`
+	// Content of the thought.
+	Content string `json:"content"`
+	// Time the thought was created.
+	Created time.Time `json:"created"`
+	// Whether the thought has been revised.
+	Revised bool `json:"revised"`
+	// Index of parent thought, or nil if this is a root for branching.
+	ParentIndex *int `json:"parentIndex,omitempty"`
+}
+
+// A ThinkingSession is an active thinking session.
+type ThinkingSession struct {
+	// Globally unique ID of the session.
+	ID string `json:"id"`
+	// Problem to solve.
+	Problem string `json:"problem"`
+	// Thoughts in the session.
+	Thoughts []*Thought `json:"thoughts"`
+	// Current thought index.
+	CurrentThought int `json:"currentThought"`
+	// Estimated total number of thoughts.
+	EstimatedTotal int `json:"estimatedTotal"`
+	// Status of the session.
+	Status string `json:"status"` // "active", "completed", "paused"
+	// Time the session was created.
+	Created time.Time `json:"created"`
+	// Time the session was last active.
+	LastActivity time.Time `json:"lastActivity"`
+	// Branches in the session. Alternative thought paths.
+	Branches []string `json:"branches,omitempty"`
+	// Version for optimistic concurrency control.
+	Version int `json:"version"`
+}
+
+// clone returns a deep copy of the ThinkingSession.
+func (s *ThinkingSession) clone() *ThinkingSession {
+	sessionCopy := *s
+	sessionCopy.Thoughts = deepCopyThoughts(s.Thoughts)
+	sessionCopy.Branches = slices.Clone(s.Branches)
+	return &sessionCopy
+}
+
+// A SessionStore is a global session store (in a real implementation, this might be a database).
+//
+// Locking Strategy:
+// The SessionStore uses a RWMutex to protect the sessions map from concurrent access.
+// All ThinkingSession modifications happen on deep copies, never on shared instances.
+// This means:
+// - Read locks protect map access.
+// - Write locks protect map modifications (adding/removing/replacing sessions)
+// - Session field modifications always happen on local copies via CompareAndSwap
+// - No shared ThinkingSession state is ever modified directly
+type SessionStore struct {
+	mu       sync.RWMutex
+	sessions map[string]*ThinkingSession // key is session ID
+}
+
+// NewSessionStore creates a new session store for managing thinking sessions.
+func NewSessionStore() *SessionStore {
+	return &SessionStore{
+		sessions: make(map[string]*ThinkingSession),
+	}
+}
+
+// Session retrieves a thinking session by ID, returning the session and whether it exists.
+func (s *SessionStore) Session(id string) (*ThinkingSession, bool) {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+	session, exists := s.sessions[id]
+	return session, exists
+}
+
+// SetSession stores or updates a thinking session in the store.
+func (s *SessionStore) SetSession(session *ThinkingSession) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.sessions[session.ID] = session
+}
+
+// CompareAndSwap atomically updates a session if the version matches.
+// Returns true if the update succeeded, false if there was a version mismatch.
+//
+// This method implements optimistic concurrency control:
+// 1. Read lock to safely access the map and copy the session
+// 2. Deep copy the session (all modifications happen on this copy)
+// 3. Release read lock and apply updates to the copy
+// 4. Write lock to check version and atomically update if unchanged
+//
+// The read lock in step 1 is necessary to prevent map access races,
+// not to protect ThinkingSession fields (which are never modified in-place).
+func (s *SessionStore) CompareAndSwap(sessionID string, updateFunc func(*ThinkingSession) (*ThinkingSession, error)) error {
+	for {
+		// Get current session
+		s.mu.RLock()
+		current, exists := s.sessions[sessionID]
+		if !exists {
+			s.mu.RUnlock()
+			return fmt.Errorf("session %s not found", sessionID)
+		}
+		// Create a deep copy
+		sessionCopy := current.clone()
+		oldVersion := current.Version
+		s.mu.RUnlock()
+
+		// Apply the update
+		updated, err := updateFunc(sessionCopy)
+		if err != nil {
+			return err
+		}
+
+		// Try to save
+		s.mu.Lock()
+		current, exists = s.sessions[sessionID]
+		if !exists {
+			s.mu.Unlock()
+			return fmt.Errorf("session %s not found", sessionID)
+		}
+		if current.Version != oldVersion {
+			// Version mismatch, retry
+			s.mu.Unlock()
+			continue
+		}
+		updated.Version = oldVersion + 1
+		s.sessions[sessionID] = updated
+		s.mu.Unlock()
+		return nil
+	}
+}
+
+// Sessions returns all thinking sessions in the store.
+func (s *SessionStore) Sessions() []*ThinkingSession {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+	return slices.Collect(maps.Values(s.sessions))
+}
+
+// SessionsSnapshot returns a deep copy of all sessions for safe concurrent access.
+func (s *SessionStore) SessionsSnapshot() []*ThinkingSession {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+
+	var sessions []*ThinkingSession
+	for _, session := range s.sessions {
+		sessions = append(sessions, session.clone())
+	}
+	return sessions
+}
+
+// SessionSnapshot returns a deep copy of a session for safe concurrent access.
+// The second return value reports whether a session with the given id exists.
+func (s *SessionStore) SessionSnapshot(id string) (*ThinkingSession, bool) {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+
+	session, exists := s.sessions[id]
+	if !exists {
+		return nil, false
+	}
+
+	return session.clone(), true
+}
+
+var store = NewSessionStore()
+
+// StartThinkingArgs are the arguments for starting a new thinking session.
+type StartThinkingArgs struct {
+	Problem        string `json:"problem"`
+	SessionID      string `json:"sessionId,omitempty"`
+	EstimatedSteps int    `json:"estimatedSteps,omitempty"`
+}
+
+// ContinueThinkingArgs are the arguments for continuing a thinking session.
+type ContinueThinkingArgs struct {
+	SessionID      string `json:"sessionId"`
+	Thought        string `json:"thought"`
+	NextNeeded     *bool  `json:"nextNeeded,omitempty"`
+	ReviseStep     *int   `json:"reviseStep,omitempty"`
+	CreateBranch   bool   `json:"createBranch,omitempty"`
+	EstimatedTotal int    `json:"estimatedTotal,omitempty"`
+}
+
+// ReviewThinkingArgs are the arguments for reviewing a thinking session.
+type ReviewThinkingArgs struct {
+	SessionID string `json:"sessionId"`
+}
+
+// ThinkingHistoryArgs are the arguments for retrieving thinking history.
+type ThinkingHistoryArgs struct {
+	SessionID string `json:"sessionId"`
+}
+
+// deepCopyThoughts creates a deep copy of a slice of thoughts.
+func deepCopyThoughts(thoughts []*Thought) []*Thought {
+	thoughtsCopy := make([]*Thought, len(thoughts))
+	for i, t := range thoughts {
+		t2 := *t
+		thoughtsCopy[i] = &t2
+	}
+	return thoughtsCopy
+}
+
+// StartThinking begins a new sequential thinking session for a complex problem.
+func StartThinking(ctx context.Context, ss *mcp.ServerSession, params *mcp.CallToolParamsFor[StartThinkingArgs]) (*mcp.CallToolResultFor[any], error) {
+	args := params.Arguments
+
+	sessionID := args.SessionID
+	if sessionID == "" {
+		sessionID = randText()
+	}
+
+	estimatedSteps := args.EstimatedSteps
+	if estimatedSteps == 0 {
+		estimatedSteps = 5 // Default estimate
+	}
+
+	session := &ThinkingSession{
+		ID:             sessionID,
+		Problem:        args.Problem,
+		EstimatedTotal: estimatedSteps,
+		Status:         "active",
+		Created:        time.Now(),
+		LastActivity:   time.Now(),
+	}
+
+	store.SetSession(session)
+
+	return &mcp.CallToolResultFor[any]{
+		Content: []mcp.Content{
+			&mcp.TextContent{
+				Text: fmt.Sprintf("Started thinking session '%s' for problem: %s\nEstimated steps: %d\nReady for your first thought.",
+					sessionID, args.Problem, estimatedSteps),
+			},
+		},
+	}, nil
+}
+
+// ContinueThinking adds the next thought step, revises a previous step, or creates a branch in the thinking process.
+func ContinueThinking(ctx context.Context, ss *mcp.ServerSession, params *mcp.CallToolParamsFor[ContinueThinkingArgs]) (*mcp.CallToolResultFor[any], error) {
+	args := params.Arguments
+
+	// Handle revision of existing thought
+	if args.ReviseStep != nil {
+		err := store.CompareAndSwap(args.SessionID, func(session *ThinkingSession) (*ThinkingSession, error) {
+			stepIndex := *args.ReviseStep - 1
+			if stepIndex < 0 || stepIndex >= len(session.Thoughts) {
+				return nil, fmt.Errorf("invalid step number: %d", *args.ReviseStep)
+			}
+
+			session.Thoughts[stepIndex].Content = args.Thought
+			session.Thoughts[stepIndex].Revised = true
+			session.LastActivity = time.Now()
+			return session, nil
+		})
+		if err != nil {
+			return nil, err
+		}
+
+		return &mcp.CallToolResultFor[any]{
+			Content: []mcp.Content{
+				&mcp.TextContent{
+					Text: fmt.Sprintf("Revised step %d in session '%s':\n%s",
+						*args.ReviseStep, args.SessionID, args.Thought),
+				},
+			},
+		}, nil
+	}
+
+	// Handle branching
+	if args.CreateBranch {
+		var branchID string
+		var branchSession *ThinkingSession
+
+		err := store.CompareAndSwap(args.SessionID, func(session *ThinkingSession) (*ThinkingSession, error) {
+			branchID = fmt.Sprintf("%s_branch_%d", args.SessionID, len(session.Branches)+1)
+			session.Branches = append(session.Branches, branchID)
+			session.LastActivity = time.Now()
+
+			// Create a new session for the branch (deep copy thoughts)
+			thoughtsCopy := deepCopyThoughts(session.Thoughts)
+			branchSession = &ThinkingSession{
+				ID:             branchID,
+				Problem:        session.Problem + " (Alternative branch)",
+				Thoughts:       thoughtsCopy,
+				CurrentThought: len(session.Thoughts),
+				EstimatedTotal: session.EstimatedTotal,
+				Status:         "active",
+				Created:        time.Now(),
+				LastActivity:   time.Now(),
+			}
+
+			return session, nil
+		})
+		if err != nil {
+			return nil, err
+		}
+
+		// Save the branch session
+		store.SetSession(branchSession)
+
+		return &mcp.CallToolResultFor[any]{
+			Content: []mcp.Content{
+				&mcp.TextContent{
+					Text: fmt.Sprintf("Created branch '%s' from session '%s'. You can now continue thinking in either session.",
+						branchID, args.SessionID),
+				},
+			},
+		}, nil
+	}
+
+	// Add new thought
+	var thoughtID int
+	var progress string
+	var statusMsg string
+
+	err := store.CompareAndSwap(args.SessionID, func(session *ThinkingSession) (*ThinkingSession, error) {
+		thoughtID = len(session.Thoughts) + 1
+		thought := &Thought{
+			Index:   thoughtID,
+			Content: args.Thought,
+			Created: time.Now(),
+			Revised: false,
+		}
+
+		session.Thoughts = append(session.Thoughts, thought)
+		session.CurrentThought = thoughtID
+		session.LastActivity = time.Now()
+
+		// Update estimated total if provided
+		if args.EstimatedTotal > 0 {
+			session.EstimatedTotal = args.EstimatedTotal
+		}
+
+		// Check if thinking is complete
+		if args.NextNeeded != nil && !*args.NextNeeded {
+			session.Status = "completed"
+		}
+
+		// Prepare response strings
+		progress = fmt.Sprintf("Step %d", thoughtID)
+		if session.EstimatedTotal > 0 {
+			progress += fmt.Sprintf(" of ~%d", session.EstimatedTotal)
+		}
+
+		if session.Status == "completed" {
+			statusMsg = "\n✓ Thinking process completed!"
+		} else {
+			statusMsg = "\nReady for next thought..."
+		}
+
+		return session, nil
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	return &mcp.CallToolResultFor[any]{
+		Content: []mcp.Content{
+			&mcp.TextContent{
+				Text: fmt.Sprintf("Session '%s' - %s:\n%s%s",
+					args.SessionID, progress, args.Thought, statusMsg),
+			},
+		},
+	}, nil
+}
+
+// ReviewThinking provides a complete review of the thinking process for a session.
+func ReviewThinking(ctx context.Context, ss *mcp.ServerSession, params *mcp.CallToolParamsFor[ReviewThinkingArgs]) (*mcp.CallToolResultFor[any], error) {
+	args := params.Arguments
+
+	// Get a snapshot of the session to avoid race conditions
+	sessionSnapshot, exists := store.SessionSnapshot(args.SessionID)
+	if !exists {
+		return nil, fmt.Errorf("session %s not found", args.SessionID)
+	}
+
+	var review strings.Builder
+	fmt.Fprintf(&review, "=== Thinking Review: %s ===\n", sessionSnapshot.ID)
+	fmt.Fprintf(&review, "Problem: %s\n", sessionSnapshot.Problem)
+	fmt.Fprintf(&review, "Status: %s\n", sessionSnapshot.Status)
+	fmt.Fprintf(&review, "Steps: %d of ~%d\n", len(sessionSnapshot.Thoughts), sessionSnapshot.EstimatedTotal)
+
+	if len(sessionSnapshot.Branches) > 0 {
+		fmt.Fprintf(&review, "Branches: %s\n", strings.Join(sessionSnapshot.Branches, ", "))
+	}
+
+	fmt.Fprintf(&review, "\n--- Thought Sequence ---\n")
+
+	for i, thought := range sessionSnapshot.Thoughts {
+		status := ""
+		if thought.Revised {
+			status = " (revised)"
+		}
+		fmt.Fprintf(&review, "%d. %s%s\n", i+1, thought.Content, status)
+	}
+
+	return &mcp.CallToolResultFor[any]{
+		Content: []mcp.Content{
+			&mcp.TextContent{
+				Text: review.String(),
+			},
+		},
+	}, nil
+}
+
+// ThinkingHistory handles resource requests for thinking session data and history.
+func ThinkingHistory(ctx context.Context, ss *mcp.ServerSession, params *mcp.ReadResourceParams) (*mcp.ReadResourceResult, error) {
+	// Extract session ID from URI (e.g., "thinking://session_123")
+	u, err := url.Parse(params.URI)
+	if err != nil {
+		return nil, fmt.Errorf("invalid thinking resource URI: %s", params.URI)
+	}
+	if u.Scheme != "thinking" {
+		return nil, fmt.Errorf("invalid thinking resource URI scheme: %s", u.Scheme)
+	}
+
+	sessionID := u.Host
+	if sessionID == "sessions" {
+		// List all sessions - use snapshot for thread safety
+		sessions := store.SessionsSnapshot()
+		data, err := json.MarshalIndent(sessions, "", "  ")
+		if err != nil {
+			return nil, fmt.Errorf("failed to marshal sessions: %w", err)
+		}
+
+		return &mcp.ReadResourceResult{
+			Contents: []*mcp.ResourceContents{
+				{
+					URI:      params.URI,
+					MIMEType: "application/json",
+					Text:     string(data),
+				},
+			},
+		}, nil
+	}
+
+	// Get specific session - use snapshot for thread safety
+	session, exists := store.SessionSnapshot(sessionID)
+	if !exists {
+		return nil, fmt.Errorf("session %s not found", sessionID)
+	}
+
+	data, err := json.MarshalIndent(session, "", "  ")
+	if err != nil {
+		return nil, fmt.Errorf("failed to marshal session: %w", err)
+	}
+
+	return &mcp.ReadResourceResult{
+		Contents: []*mcp.ResourceContents{
+			{
+				URI:      params.URI,
+				MIMEType: "application/json",
+				Text:     string(data),
+			},
+		},
+	}, nil
+}
+
+// Copied from crypto/rand.
+// TODO: once 1.24 is assured, just use crypto/rand.
+const base32alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZ234567"
+
+func randText() string {
+	// ⌈log₃₂ 2¹²⁸⌉ = 26 chars
+	src := make([]byte, 26)
+	rand.Read(src)
+	for i := range src {
+		src[i] = base32alphabet[src[i]%32]
+	}
+	return string(src)
+}
+
+func main() {
+	flag.Parse()
+
+	server := mcp.NewServer(&mcp.Implementation{Name: "sequential-thinking"}, nil)
+
+	mcp.AddTool(server, &mcp.Tool{
+		Name:        "start_thinking",
+		Description: "Begin a new sequential thinking session for a complex problem",
+	}, StartThinking)
+
+	mcp.AddTool(server, &mcp.Tool{
+		Name:        "continue_thinking",
+		Description: "Add the next thought step, revise a previous step, or create a branch",
+	}, ContinueThinking)
+
+	mcp.AddTool(server, &mcp.Tool{
+		Name:        "review_thinking",
+		Description: "Review the complete thinking process for a session",
+	}, ReviewThinking)
+
+	server.AddResource(&mcp.Resource{
+		Name:        "thinking_sessions",
+		Description: "Access thinking session data and history",
+		URI:         "thinking://sessions",
+		MIMEType:    "application/json",
+	}, ThinkingHistory)
+
+	if *httpAddr != "" {
+		handler := mcp.NewStreamableHTTPHandler(func(*http.Request) *mcp.Server {
+			return server
+		}, nil)
+		log.Printf("Sequential Thinking MCP server listening at %s", *httpAddr)
+		if err := http.ListenAndServe(*httpAddr, handler); err != nil {
+			log.Fatal(err)
+		}
+	} else {
+		t := mcp.NewLoggingTransport(mcp.NewStdioTransport(), os.Stderr)
+		if err := server.Run(context.Background(), t); err != nil {
+			log.Printf("Server failed: %v", err)
+		}
+	}
+}

examples/sequentialthinking/main_test.go

@@ -0,0 +1,548 @@
+// Copyright 2025 The Go MCP SDK Authors. All rights reserved.
+// Use of this source code is governed by an MIT-style
+// license that can be found in the LICENSE file.
+
+package main
+
+import (
+	"context"
+	"encoding/json"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/modelcontextprotocol/go-sdk/mcp"
+)
+
+func TestStartThinking(t *testing.T) {
+	// Reset store for clean test
+	store = NewSessionStore()
+
+	ctx := context.Background()
+
+	args := StartThinkingArgs{
+		Problem:        "How to implement a binary search algorithm",
+		SessionID:      "test_session",
+		EstimatedSteps: 5,
+	}
+
+	params := &mcp.CallToolParamsFor[StartThinkingArgs]{
+		Name:      "start_thinking",
+		Arguments: args,
+	}
+
+	result, err := StartThinking(ctx, nil, params)
+	if err != nil {
+		t.Fatalf("StartThinking() error = %v", err)
+	}
+
+	if len(result.Content) == 0 {
+		t.Fatal("No content in result")
+	}
+
+	textContent, ok := result.Content[0].(*mcp.TextContent)
+	if !ok {
+		t.Fatal("Expected TextContent")
+	}
+
+	if !strings.Contains(textContent.Text, "test_session") {
+		t.Error("Result should contain session ID")
+	}
+
+	if !strings.Contains(textContent.Text, "How to implement a binary search algorithm") {
+		t.Error("Result should contain the problem statement")
+	}
+
+	// Verify session was stored
+	session, exists := store.Session("test_session")
+	if !exists {
+		t.Fatal("Session was not stored")
+	}
+
+	if session.Problem != args.Problem {
+		t.Errorf("Expected problem %s, got %s", args.Problem, session.Problem)
+	}
+
+	if session.EstimatedTotal != 5 {
+		t.Errorf("Expected estimated total 5, got %d", session.EstimatedTotal)
+	}
+
+	if session.Status != "active" {
+		t.Errorf("Expected status 'active', got %s", session.Status)
+	}
+}
+
+func TestContinueThinking(t *testing.T) {
+	// Reset store and create initial session
+	store = NewSessionStore()
+
+	// First start a thinking session
+	ctx := context.Background()
+	startArgs := StartThinkingArgs{
+		Problem:        "Test problem",
+		SessionID:      "test_continue",
+		EstimatedSteps: 3,
+	}
+
+	startParams := &mcp.CallToolParamsFor[StartThinkingArgs]{
+		Name:      "start_thinking",
+		Arguments: startArgs,
+	}
+
+	_, err := StartThinking(ctx, nil, startParams)
+	if err != nil {
+		t.Fatalf("StartThinking() error = %v", err)
+	}
+
+	// Now continue thinking
+	continueArgs := ContinueThinkingArgs{
+		SessionID: "test_continue",
+		Thought:   "First thought: I need to understand the problem",
+	}
+
+	continueParams := &mcp.CallToolParamsFor[ContinueThinkingArgs]{
+		Name:      "continue_thinking",
+		Arguments: continueArgs,
+	}
+
+	result, err := ContinueThinking(ctx, nil, continueParams)
+	if err != nil {
+		t.Fatalf("ContinueThinking() error = %v", err)
+	}
+
+	// Verify result
+	if len(result.Content) == 0 {
+		t.Fatal("No content in result")
+	}
+
+	textContent, ok := result.Content[0].(*mcp.TextContent)
+	if !ok {
+		t.Fatal("Expected TextContent")
+	}
+
+	if !strings.Contains(textContent.Text, "Step 1") {
+		t.Error("Result should contain step number")
+	}
+
+	// Verify session was updated
+	session, exists := store.Session("test_continue")
+	if !exists {
+		t.Fatal("Session not found")
+	}
+
+	if len(session.Thoughts) != 1 {
+		t.Errorf("Expected 1 thought, got %d", len(session.Thoughts))
+	}
+
+	if session.Thoughts[0].Content != continueArgs.Thought {
+		t.Errorf("Expected thought content %s, got %s", continueArgs.Thought, session.Thoughts[0].Content)
+	}
+
+	if session.CurrentThought != 1 {
+		t.Errorf("Expected current thought 1, got %d", session.CurrentThought)
+	}
+}
+
+func TestContinueThinkingWithCompletion(t *testing.T) {
+	// Reset store and create initial session
+	store = NewSessionStore()
+
+	ctx := context.Background()
+	startArgs := StartThinkingArgs{
+		Problem:   "Simple test",
+		SessionID: "test_completion",
+	}
+
+	startParams := &mcp.CallToolParamsFor[StartThinkingArgs]{
+		Name:      "start_thinking",
+		Arguments: startArgs,
+	}
+
+	_, err := StartThinking(ctx, nil, startParams)
+	if err != nil {
+		t.Fatalf("StartThinking() error = %v", err)
+	}
+
+	// Continue with completion flag
+	nextNeeded := false
+	continueArgs := ContinueThinkingArgs{
+		SessionID:  "test_completion",
+		Thought:    "Final thought",
+		NextNeeded: &nextNeeded,
+	}
+
+	continueParams := &mcp.CallToolParamsFor[ContinueThinkingArgs]{
+		Name:      "continue_thinking",
+		Arguments: continueArgs,
+	}
+
+	result, err := ContinueThinking(ctx, nil, continueParams)
+	if err != nil {
+		t.Fatalf("ContinueThinking() error = %v", err)
+	}
+
+	// Check completion message
+	textContent, ok := result.Content[0].(*mcp.TextContent)
+	if !ok {
+		t.Fatal("Expected TextContent")
+	}
+
+	if !strings.Contains(textContent.Text, "completed") {
+		t.Error("Result should indicate completion")
+	}
+
+	// Verify session status
+	session, exists := store.Session("test_completion")
+	if !exists {
+		t.Fatal("Session not found")
+	}
+
+	if session.Status != "completed" {
+		t.Errorf("Expected status 'completed', got %s", session.Status)
+	}
+}
+
+func TestContinueThinkingRevision(t *testing.T) {
+	// Setup session with existing thoughts
+	store = NewSessionStore()
+	session := &ThinkingSession{
+		ID:      "test_revision",
+		Problem: "Test problem",
+		Thoughts: []*Thought{
+			{Index: 1, Content: "Original thought", Created: time.Now()},
+			{Index: 2, Content: "Second thought", Created: time.Now()},
+		},
+		CurrentThought: 2,
+		EstimatedTotal: 3,
+		Status:         "active",
+		Created:        time.Now(),
+		LastActivity:   time.Now(),
+	}
+	store.SetSession(session)
+
+	ctx := context.Background()
+	reviseStep := 1
+	continueArgs := ContinueThinkingArgs{
+		SessionID:  "test_revision",
+		Thought:    "Revised first thought",
+		ReviseStep: &reviseStep,
+	}
+
+	continueParams := &mcp.CallToolParamsFor[ContinueThinkingArgs]{
+		Name:      "continue_thinking",
+		Arguments: continueArgs,
+	}
+
+	result, err := ContinueThinking(ctx, nil, continueParams)
+	if err != nil {
+		t.Fatalf("ContinueThinking() error = %v", err)
+	}
+
+	// Verify revision message
+	textContent, ok := result.Content[0].(*mcp.TextContent)
+	if !ok {
+		t.Fatal("Expected TextContent")
+	}
+
+	if !strings.Contains(textContent.Text, "Revised step 1") {
+		t.Error("Result should indicate revision")
+	}
+
+	// Verify thought was revised
+	updatedSession, _ := store.Session("test_revision")
+	if updatedSession.Thoughts[0].Content != "Revised first thought" {
+		t.Error("First thought should be revised")
+	}
+
+	if !updatedSession.Thoughts[0].Revised {
+		t.Error("First thought should be marked as revised")
+	}
+}
+
+func TestContinueThinkingBranching(t *testing.T) {
+	// Setup session with existing thoughts
+	store = NewSessionStore()
+	session := &ThinkingSession{
+		ID:      "test_branch",
+		Problem: "Test problem",
+		Thoughts: []*Thought{
+			{Index: 1, Content: "First thought", Created: time.Now()},
+		},
+		CurrentThought: 1,
+		EstimatedTotal: 3,
+		Status:         "active",
+		Created:        time.Now(),
+		LastActivity:   time.Now(),
+		Branches:       []string{},
+	}
+	store.SetSession(session)
+
+	ctx := context.Background()
+	continueArgs := ContinueThinkingArgs{
+		SessionID:    "test_branch",
+		Thought:      "Alternative approach",
+		CreateBranch: true,
+	}
+
+	continueParams := &mcp.CallToolParamsFor[ContinueThinkingArgs]{
+		Name:      "continue_thinking",
+		Arguments: continueArgs,
+	}
+
+	result, err := ContinueThinking(ctx, nil, continueParams)
+	if err != nil {
+		t.Fatalf("ContinueThinking() error = %v", err)
+	}
+
+	// Verify branch creation message
+	textContent, ok := result.Content[0].(*mcp.TextContent)
+	if !ok {
+		t.Fatal("Expected TextContent")
+	}
+
+	if !strings.Contains(textContent.Text, "Created branch") {
+		t.Error("Result should indicate branch creation")
+	}
+
+	// Verify branch was created
+	updatedSession, _ := store.Session("test_branch")
+	if len(updatedSession.Branches) != 1 {
+		t.Errorf("Expected 1 branch, got %d", len(updatedSession.Branches))
+	}
+
+	branchID := updatedSession.Branches[0]
+	if !strings.Contains(branchID, "test_branch_branch_") {
+		t.Error("Branch ID should contain parent session ID")
+	}
+
+	// Verify branch session exists
+	branchSession, exists := store.Session(branchID)
+	if !exists {
+		t.Fatal("Branch session should exist")
+	}
+
+	if len(branchSession.Thoughts) != 1 {
+		t.Error("Branch should inherit parent thoughts")
+	}
+}
+
+func TestReviewThinking(t *testing.T) {
+	// Setup session with thoughts
+	store = NewSessionStore()
+	session := &ThinkingSession{
+		ID:      "test_review",
+		Problem: "Complex problem",
+		Thoughts: []*Thought{
+			{Index: 1, Content: "First thought", Created: time.Now(), Revised: false},
+			{Index: 2, Content: "Second thought", Created: time.Now(), Revised: true},
+			{Index: 3, Content: "Final thought", Created: time.Now(), Revised: false},
+		},
+		CurrentThought: 3,
+		EstimatedTotal: 3,
+		Status:         "completed",
+		Created:        time.Now(),
+		LastActivity:   time.Now(),
+		Branches:       []string{"test_review_branch_1"},
+	}
+	store.SetSession(session)
+
+	ctx := context.Background()
+	reviewArgs := ReviewThinkingArgs{
+		SessionID: "test_review",
+	}
+
+	reviewParams := &mcp.CallToolParamsFor[ReviewThinkingArgs]{
+		Name:      "review_thinking",
+		Arguments: reviewArgs,
+	}
+
+	result, err := ReviewThinking(ctx, nil, reviewParams)
+	if err != nil {
+		t.Fatalf("ReviewThinking() error = %v", err)
+	}
+
+	// Verify review content
+	textContent, ok := result.Content[0].(*mcp.TextContent)
+	if !ok {
+		t.Fatal("Expected TextContent")
+	}
+
+	reviewText := textContent.Text
+
+	if !strings.Contains(reviewText, "test_review") {
+		t.Error("Review should contain session ID")
+	}
+
+	if !strings.Contains(reviewText, "Complex problem") {
+		t.Error("Review should contain problem")
+	}
+
+	if !strings.Contains(reviewText, "completed") {
+		t.Error("Review should contain status")
+	}
+
+	if !strings.Contains(reviewText, "Steps: 3 of ~3") {
+		t.Error("Review should contain step count")
+	}
+
+	if !strings.Contains(reviewText, "First thought") {
+		t.Error("Review should contain first thought")
+	}
+
+	if !strings.Contains(reviewText, "(revised)") {
+		t.Error("Review should indicate revised thoughts")
+	}
+
+	if !strings.Contains(reviewText, "test_review_branch_1") {
+		t.Error("Review should list branches")
+	}
+}
+
+func TestThinkingHistory(t *testing.T) {
+	// Setup test sessions
+	store = NewSessionStore()
+	session1 := &ThinkingSession{
+		ID:             "session1",
+		Problem:        "Problem 1",
+		Thoughts:       []*Thought{{Index: 1, Content: "Thought 1", Created: time.Now()}},
+		CurrentThought: 1,
+		EstimatedTotal: 2,
+		Status:         "active",
+		Created:        time.Now(),
+		LastActivity:   time.Now(),
+	}
+	session2 := &ThinkingSession{
+		ID:             "session2",
+		Problem:        "Problem 2",
+		Thoughts:       []*Thought{{Index: 1, Content: "Thought 1", Created: time.Now()}},
+		CurrentThought: 1,
+		EstimatedTotal: 3,
+		Status:         "completed",
+		Created:        time.Now(),
+		LastActivity:   time.Now(),
+	}
+	store.SetSession(session1)
+	store.SetSession(session2)
+
+	ctx := context.Background()
+
+	// Test listing all sessions
+	listParams := &mcp.ReadResourceParams{
+		URI: "thinking://sessions",
+	}
+
+	result, err := ThinkingHistory(ctx, nil, listParams)
+	if err != nil {
+		t.Fatalf("ThinkingHistory() error = %v", err)
+	}
+
+	if len(result.Contents) != 1 {
+		t.Fatal("Expected 1 content item")
+	}
+
+	content := result.Contents[0]
+	if content.MIMEType != "application/json" {
+		t.Error("Expected JSON MIME type")
+	}
+
+	// Parse and verify sessions list
+	var sessions []*ThinkingSession
+	err = json.Unmarshal([]byte(content.Text), &sessions)
+	if err != nil {
+		t.Fatalf("Failed to parse sessions JSON: %v", err)
+	}
+
+	if len(sessions) != 2 {
+		t.Errorf("Expected 2 sessions, got %d", len(sessions))
+	}
+
+	// Test getting specific session
+	sessionParams := &mcp.ReadResourceParams{
+		URI: "thinking://session1",
+	}
+
+	result, err = ThinkingHistory(ctx, nil, sessionParams)
+	if err != nil {
+		t.Fatalf("ThinkingHistory() error = %v", err)
+	}
+
+	var retrievedSession ThinkingSession
+	err = json.Unmarshal([]byte(result.Contents[0].Text), &retrievedSession)
+	if err != nil {
+		t.Fatalf("Failed to parse session JSON: %v", err)
+	}
+
+	if retrievedSession.ID != "session1" {
+		t.Errorf("Expected session ID 'session1', got %s", retrievedSession.ID)
+	}
+
+	if retrievedSession.Problem != "Problem 1" {
+		t.Errorf("Expected problem 'Problem 1', got %s", retrievedSession.Problem)
+	}
+}
+
+func TestInvalidOperations(t *testing.T) {
+	store = NewSessionStore()
+	ctx := context.Background()
+
+	// Test continue thinking with non-existent session
+	continueArgs := ContinueThinkingArgs{
+		SessionID: "nonexistent",
+		Thought:   "Some thought",
+	}
+
+	continueParams := &mcp.CallToolParamsFor[ContinueThinkingArgs]{
+		Name:      "continue_thinking",
+		Arguments: continueArgs,
+	}
+
+	_, err := ContinueThinking(ctx, nil, continueParams)
+	if err == nil {
+		t.Error("Expected error for non-existent session")
+	}
+
+	// Test review with non-existent session
+	reviewArgs := ReviewThinkingArgs{
+		SessionID: "nonexistent",
+	}
+
+	reviewParams := &mcp.CallToolParamsFor[ReviewThinkingArgs]{
+		Name:      "review_thinking",
+		Arguments: reviewArgs,
+	}
+
+	_, err = ReviewThinking(ctx, nil, reviewParams)
+	if err == nil {
+		t.Error("Expected error for non-existent session in review")
+	}
+
+	// Test invalid revision step
+	session := &ThinkingSession{
+		ID:             "test_invalid",
+		Problem:        "Test",
+		Thoughts:       []*Thought{{Index: 1, Content: "Thought", Created: time.Now()}},
+		CurrentThought: 1,
+		EstimatedTotal: 2,
+		Status:         "active",
+		Created:        time.Now(),
+		LastActivity:   time.Now(),
+	}
+	store.SetSession(session)
+
+	reviseStep := 5 // Invalid step number
+	invalidReviseArgs := ContinueThinkingArgs{
+		SessionID:  "test_invalid",
+		Thought:    "Revised",
+		ReviseStep: &reviseStep,
+	}
+
+	invalidReviseParams := &mcp.CallToolParamsFor[ContinueThinkingArgs]{
+		Name:      "continue_thinking",
+		Arguments: invalidReviseArgs,
+	}
+
+	_, err = ContinueThinking(ctx, nil, invalidReviseParams)
+	if err == nil {
+		t.Error("Expected error for invalid revision step")
+	}
+}