From ed43a61bf88f3238d4844391e1f7926d8490ae50 Mon Sep 17 00:00:00 2001 From: pdfinn Date: Wed, 4 Feb 2026 12:14:53 +0700 Subject: [PATCH] feat(llm9p): Add per-fid session isolation and prefill support - Add SessionManager for per-fid conversation isolation - Each 9P fid now gets its own conversation history - Add FidAwareFile interface for files needing fid context - Add /n/llm/prefill file for assistant response prefill - Prefill helps keep model in character (e.g., "[Veltro]") - Update ask, new, context files to use session manager - Fix context contamination between parent and subagent Co-Authored-By: Claude Opus 4.5 --- cmd/llm9p/main.go | 5 +- internal/llm/backend.go | 8 ++ internal/llm/cli_client.go | 106 +++++++++++++++ internal/llm/client.go | 142 +++++++++++++++++++- internal/llm/session.go | 192 ++++++++++++++++++++++++++++ internal/llmfs/ask.go | 82 ++++++------ internal/llmfs/ask_test.go | 178 +++++++++++++------------- internal/llmfs/context.go | 49 +++++-- internal/llmfs/mock_backend_test.go | 11 ++ internal/llmfs/new.go | 33 ++++- internal/llmfs/prefill.go | 54 ++++++++ internal/llmfs/root.go | 43 ++++--- internal/protocol/fs.go | 16 +++ internal/protocol/server.go | 25 +++- 14 files changed, 781 insertions(+), 163 deletions(-) create mode 100644 internal/llm/session.go create mode 100644 internal/llmfs/prefill.go diff --git a/cmd/llm9p/main.go b/cmd/llm9p/main.go index 5431ba6be0f1656c91a7d59f313d9b6577da9637..d9f5718d6f16d63c95f766619453b8cdf9e91988 100644 --- a/cmd/llm9p/main.go +++ b/cmd/llm9p/main.go @@ -69,8 +69,11 @@ func main() { os.Exit(1) } + // Create session manager for per-fid isolation + sm := llm.NewSessionManager(client) + // Create filesystem - root := llmfs.NewRoot(client) + root := llmfs.NewRoot(sm) // Create 9P server server := protocol.NewServer(root) diff --git a/internal/llm/backend.go b/internal/llm/backend.go index 8fa0715af7269532fc2a267880601447be5c64ac..30b13cf19197c58f5f4ada8450115084aad4c590 100644 --- a/internal/llm/backend.go +++ b/internal/llm/backend.go @@ -22,6 +22,11 @@ type Backend interface { ThinkingTokens() int // SetThinkingTokens sets the thinking token budget SetThinkingTokens(tokens int) + // Prefill returns the assistant response prefill string + Prefill() string + // SetPrefill sets a string to prefill the assistant response + // This helps keep the model in character (e.g., "[Veltro] ") + SetPrefill(prefill string) // LastTokens returns token count from last response LastTokens() int // TotalTokens returns cumulative token count for this conversation @@ -41,6 +46,9 @@ type Backend interface { Reset() // Ask sends a prompt and returns the response (blocking) Ask(ctx context.Context, prompt string) (string, error) + // AskWithHistory sends a prompt with explicit message history (for per-fid isolation) + // Returns response text and token count + AskWithHistory(ctx context.Context, history []Message, prompt string) (string, int, error) // StartStream begins streaming a response StartStream(ctx context.Context, prompt string) error // ReadStreamChunk reads the next streaming chunk diff --git a/internal/llm/cli_client.go b/internal/llm/cli_client.go index e23ce9fc863bd0fb34fbcfc41db133ec9f3f3e0d..22fd7860f99112fdf44101b95aca2c9783b965be 100644 --- a/internal/llm/cli_client.go +++ b/internal/llm/cli_client.go @@ -19,6 +19,7 @@ type CLIClient struct { model string temperature float64 systemPrompt string + prefill string // assistant response prefill for keeping model in character messages []Message lastTokens int totalTokens int // cumulative estimated token count @@ -105,6 +106,20 @@ func (c *CLIClient) SetThinkingTokens(tokens int) { c.thinkingTokens = tokens } +// Prefill returns the assistant response prefill string +func (c *CLIClient) Prefill() string { + c.mu.RLock() + defer c.mu.RUnlock() + return c.prefill +} + +// SetPrefill sets a string to prefill the assistant response +func (c *CLIClient) SetPrefill(prefill string) { + c.mu.Lock() + defer c.mu.Unlock() + c.prefill = prefill +} + // SystemPrompt returns the current system prompt func (c *CLIClient) SystemPrompt() string { c.mu.RLock() @@ -549,3 +564,94 @@ func (c *CLIClient) WaitStream() { <-done } } + +// AskWithHistory sends a prompt with explicit message history for per-fid isolation. +// Unlike Ask(), this does not modify the client's internal messages state. +// Returns response text and estimated token count. +func (c *CLIClient) AskWithHistory(ctx context.Context, history []Message, prompt string) (string, int, error) { + // Get settings with lock + c.mu.RLock() + model := c.model + thinkingTokens := c.thinkingTokens + systemPromptSetting := c.systemPrompt + prefill := c.prefill + c.mu.RUnlock() + + // Build prompt from provided history + var parts []string + var systemParts []string + + // Add dedicated system prompt first + if systemPromptSetting != "" { + systemParts = append(systemParts, systemPromptSetting) + } + + for _, msg := range history { + switch msg.Role { + case "system": + systemParts = append(systemParts, msg.Content) + case "user": + parts = append(parts, fmt.Sprintf("Human: %s", msg.Content)) + case "assistant": + parts = append(parts, fmt.Sprintf("Assistant: %s", msg.Content)) + } + } + + // Add the new user prompt + parts = append(parts, fmt.Sprintf("Human: %s", prompt)) + + fullPrompt := strings.Join(parts, "\n\n") + systemPrompt := strings.Join(systemParts, "\n\n") + + // Build claude CLI command + args := []string{ + "--print", + "--output-format", "json", + "--model", model, + "--allowedTools", "", + "--dangerously-skip-permissions", + } + + if systemPrompt != "" { + args = append(args, "--system-prompt", systemPrompt) + } + + args = append(args, "-") // Read from stdin + + cmd := exec.CommandContext(ctx, "claude", args...) + cmd.Stdin = bytes.NewBufferString(fullPrompt) + + // Set thinking token budget via environment variable + cmd.Env = append(cmd.Environ(), func() string { + if thinkingTokens < 0 { + return "MAX_THINKING_TOKENS=31999" + } + return fmt.Sprintf("MAX_THINKING_TOKENS=%d", thinkingTokens) + }()) + + var stdout, stderr bytes.Buffer + cmd.Stdout = &stdout + cmd.Stderr = &stderr + + if err := cmd.Run(); err != nil { + return "", 0, fmt.Errorf("claude CLI error: %w (stderr: %s)", err, stderr.String()) + } + + // Parse JSON response + responseText, err := parseJSONResponse(stdout.String()) + if err != nil { + return "", 0, fmt.Errorf("failed to parse CLI response: %w", err) + } + + // Prepend prefill to response to keep model in character + // Note: CLI doesn't support true prefill (partial assistant message), + // so we prepend it to the response for consistent behavior with API client + if prefill != "" { + responseText = prefill + responseText + } + + // Estimate tokens: prompt + response (chars / 4) + tokens := estimateTokens(fullPrompt) + estimateTokens(responseText) + + return responseText, tokens, nil +} diff --git a/internal/llm/client.go b/internal/llm/client.go index 28d29e67af268f78e50bae5711bf60ed590fce4e..02cedced96fbc4b5a8138a454efa9040d6f682a5 100644 --- a/internal/llm/client.go +++ b/internal/llm/client.go @@ -7,6 +7,7 @@ import ( "fmt" "strings" "sync" + "time" "github.com/anthropics/anthropic-sdk-go" "github.com/anthropics/anthropic-sdk-go/option" @@ -18,6 +19,24 @@ type Message struct { Content string `json:"content"` // message content } +// MetricsCallback is called after each LLM request with performance data +type MetricsCallback func(inputTokens, outputTokens int, latencyMs int64) + +// Global metrics callback - set by llmfs to record metrics +var metricsCallback MetricsCallback + +// SetMetricsCallback registers a callback for recording metrics +func SetMetricsCallback(cb MetricsCallback) { + metricsCallback = cb +} + +// RecordMetrics calls the registered callback if set +func RecordMetrics(inputTokens, outputTokens int, latencyMs int64) { + if metricsCallback != nil { + metricsCallback(inputTokens, outputTokens, latencyMs) + } +} + // Client wraps the Anthropic API client with conversation state type Client struct { client anthropic.Client @@ -25,6 +44,7 @@ type Client struct { model string temperature float64 systemPrompt string + prefill string // assistant response prefill for keeping model in character messages []Message lastTokens int totalTokens int // cumulative token count for context tracking @@ -93,6 +113,20 @@ func (c *Client) SetThinkingTokens(tokens int) { c.thinkingTokens = tokens } +// Prefill returns the assistant response prefill string +func (c *Client) Prefill() string { + c.mu.RLock() + defer c.mu.RUnlock() + return c.prefill +} + +// SetPrefill sets a string to prefill the assistant response +func (c *Client) SetPrefill(prefill string) { + c.mu.Lock() + defer c.mu.Unlock() + c.prefill = prefill +} + // SystemPrompt returns the current system prompt func (c *Client) SystemPrompt() string { c.mu.RLock() @@ -286,8 +320,11 @@ func (c *Client) Ask(ctx context.Context, prompt string) (string, error) { params.System = systemBlocks } - // Make the API call + // Make the API call with timing + startTime := time.Now() response, err := c.client.Messages.New(ctx, params) + latencyMs := time.Since(startTime).Milliseconds() + if err != nil { // Remove the user message on error c.mu.Lock() @@ -313,6 +350,11 @@ func (c *Client) Ask(ctx context.Context, prompt string) (string, error) { c.totalTokens += c.lastTokens c.mu.Unlock() + // Record metrics (input and output tokens separately for analysis) + inputToks := int(response.Usage.InputTokens) + outputToks := int(response.Usage.OutputTokens) + RecordMetrics(inputToks, outputToks, latencyMs) + return responseText, nil } @@ -472,3 +514,101 @@ func (c *Client) WaitStream() { <-done } } + +// AskWithHistory sends a prompt with explicit message history for per-fid isolation. +// Unlike Ask(), this does not modify the client's internal messages state. +// Returns response text and token count. +func (c *Client) AskWithHistory(ctx context.Context, history []Message, prompt string) (string, int, error) { + // Get settings with lock + c.mu.RLock() + model := c.model + temp := c.temperature + systemPrompt := c.systemPrompt + prefill := c.prefill + c.mu.RUnlock() + + // Build API messages from provided history plus the new prompt + apiMessages := make([]anthropic.MessageParam, 0, len(history)+2) + var systemBlocks []anthropic.TextBlockParam + + // Add dedicated system prompt first + if systemPrompt != "" { + systemBlocks = append(systemBlocks, anthropic.TextBlockParam{ + Text: systemPrompt, + }) + } + + for _, msg := range history { + switch msg.Role { + case "system": + systemBlocks = append(systemBlocks, anthropic.TextBlockParam{ + Text: msg.Content, + }) + case "user": + apiMessages = append(apiMessages, anthropic.NewUserMessage( + anthropic.NewTextBlock(msg.Content), + )) + case "assistant": + apiMessages = append(apiMessages, anthropic.NewAssistantMessage( + anthropic.NewTextBlock(msg.Content), + )) + } + } + + // Add the new user prompt + apiMessages = append(apiMessages, anthropic.NewUserMessage( + anthropic.NewTextBlock(prompt), + )) + + // Add prefill as partial assistant message to keep model in character + // The model will continue from this point + if prefill != "" { + apiMessages = append(apiMessages, anthropic.NewAssistantMessage( + anthropic.NewTextBlock(prefill), + )) + } + + // Build request params + params := anthropic.MessageNewParams{ + Model: anthropic.Model(model), + MaxTokens: 4096, + Messages: apiMessages, + Temperature: anthropic.Float(temp), + } + + // Add system prompt if present + if len(systemBlocks) > 0 { + params.System = systemBlocks + } + + // Make the API call with timing + startTime := time.Now() + response, err := c.client.Messages.New(ctx, params) + latencyMs := time.Since(startTime).Milliseconds() + + if err != nil { + return "", 0, fmt.Errorf("API error: %w", err) + } + + // Extract response text + var responseText string + for _, block := range response.Content { + if block.Type == "text" { + responseText += block.Text + } + } + + // Prepend prefill to response (it was used as partial assistant message) + if prefill != "" { + responseText = prefill + responseText + } + + tokens := int(response.Usage.InputTokens + response.Usage.OutputTokens) + + // Record metrics + inputToks := int(response.Usage.InputTokens) + outputToks := int(response.Usage.OutputTokens) + RecordMetrics(inputToks, outputToks, latencyMs) + + return responseText, tokens, nil +} diff --git a/internal/llm/session.go b/internal/llm/session.go new file mode 100644 index 0000000000000000000000000000000000000000..4cbc38256f2268c36750a6691390bf487705a150 --- /dev/null +++ b/internal/llm/session.go @@ -0,0 +1,192 @@ +// Package llm provides LLM backends for the 9P filesystem. +package llm + +import ( + "context" + "encoding/json" + "sync" +) + +// Session holds per-fid conversation state. +// Each fid that opens the ask file gets its own session with isolated history. +type Session struct { + ID uint32 + messages []Message + lastResponse string + lastTokens int + totalTokens int + mu sync.RWMutex +} + +// NewSession creates a new session for the given fid. +func NewSession(fid uint32) *Session { + return &Session{ + ID: fid, + messages: make([]Message, 0), + } +} + +// Messages returns a copy of the session's conversation history. +func (s *Session) Messages() []Message { + s.mu.RLock() + defer s.mu.RUnlock() + result := make([]Message, len(s.messages)) + copy(result, s.messages) + return result +} + +// MessagesJSON returns the session's conversation history as JSON. +func (s *Session) MessagesJSON() ([]byte, error) { + s.mu.RLock() + defer s.mu.RUnlock() + return json.MarshalIndent(s.messages, "", " ") +} + +// AddMessage adds a message to the session's history. +func (s *Session) AddMessage(role, content string) { + s.mu.Lock() + defer s.mu.Unlock() + s.messages = append(s.messages, Message{Role: role, Content: content}) +} + +// AddSystemMessage adds a system message to the session's history. +func (s *Session) AddSystemMessage(content string) { + s.mu.Lock() + defer s.mu.Unlock() + s.messages = append([]Message{{Role: "system", Content: content}}, s.messages...) +} + +// SetLastResponse sets the last response for this session. +func (s *Session) SetLastResponse(response string) { + s.mu.Lock() + defer s.mu.Unlock() + s.lastResponse = response +} + +// LastResponse returns the last response for this session. +func (s *Session) LastResponse() string { + s.mu.RLock() + defer s.mu.RUnlock() + return s.lastResponse +} + +// LastTokens returns the token count from the last response. +func (s *Session) LastTokens() int { + s.mu.RLock() + defer s.mu.RUnlock() + return s.lastTokens +} + +// TotalTokens returns cumulative token count for this session. +func (s *Session) TotalTokens() int { + s.mu.RLock() + defer s.mu.RUnlock() + return s.totalTokens +} + +// SetTokens updates the token counts for this session. +func (s *Session) SetTokens(last, total int) { + s.mu.Lock() + defer s.mu.Unlock() + s.lastTokens = last + s.totalTokens = total +} + +// AddTokens adds to the token counts for this session. +func (s *Session) AddTokens(tokens int) { + s.mu.Lock() + defer s.mu.Unlock() + s.lastTokens = tokens + s.totalTokens += tokens +} + +// Reset clears the session's conversation history. +func (s *Session) Reset() { + s.mu.Lock() + defer s.mu.Unlock() + s.messages = make([]Message, 0) + s.lastResponse = "" + s.lastTokens = 0 + s.totalTokens = 0 +} + +// SessionManager maps fids to sessions and delegates to a shared backend. +type SessionManager struct { + sessions map[uint32]*Session + backend Backend // shared backend for API calls and global settings + mu sync.RWMutex +} + +// NewSessionManager creates a new session manager with the given backend. +func NewSessionManager(backend Backend) *SessionManager { + return &SessionManager{ + sessions: make(map[uint32]*Session), + backend: backend, + } +} + +// Backend returns the underlying shared backend. +func (sm *SessionManager) Backend() Backend { + return sm.backend +} + +// GetOrCreate returns the session for the given fid, creating one if necessary. +func (sm *SessionManager) GetOrCreate(fid uint32) *Session { + sm.mu.Lock() + defer sm.mu.Unlock() + if s, ok := sm.sessions[fid]; ok { + return s + } + s := NewSession(fid) + sm.sessions[fid] = s + return s +} + +// Get returns the session for the given fid, or nil if it doesn't exist. +func (sm *SessionManager) Get(fid uint32) *Session { + sm.mu.RLock() + defer sm.mu.RUnlock() + return sm.sessions[fid] +} + +// Remove removes the session for the given fid. +func (sm *SessionManager) Remove(fid uint32) { + sm.mu.Lock() + defer sm.mu.Unlock() + delete(sm.sessions, fid) +} + +// Reset clears the session for the given fid (but keeps the session). +func (sm *SessionManager) Reset(fid uint32) { + session := sm.GetOrCreate(fid) + session.Reset() +} + +// Ask sends a prompt using the session's conversation history. +// The response is stored in the session and returned. +func (sm *SessionManager) Ask(ctx context.Context, fid uint32, prompt string) (string, error) { + session := sm.GetOrCreate(fid) + + // Get current history before adding new message + history := session.Messages() + + // Use backend's AskWithHistory - it doesn't modify backend state + response, tokens, err := sm.backend.AskWithHistory(ctx, history, prompt) + if err != nil { + session.SetLastResponse("Error: " + err.Error()) + return "", err + } + + // Add user message and assistant response to session history + session.AddMessage("user", prompt) + session.AddMessage("assistant", response) + session.AddTokens(tokens) + session.SetLastResponse(response) + + return response, nil +} + +// ContextLimit returns the model's context window limit from the backend. +func (sm *SessionManager) ContextLimit() int { + return sm.backend.ContextLimit() +} diff --git a/internal/llmfs/ask.go b/internal/llmfs/ask.go index cc3be1b2487ea007975b3938a1f6304b217c8cd5..95c477cbb8d204fae0b54dbfc6d784f64f7318b3 100644 --- a/internal/llmfs/ask.go +++ b/internal/llmfs/ask.go @@ -5,7 +5,6 @@ import ( "io" "log" "strings" - "sync" "github.com/NERVsystems/llm9p/internal/llm" "github.com/NERVsystems/llm9p/internal/protocol" @@ -14,26 +13,38 @@ import ( // CompactThreshold is the percentage of context limit at which auto-compaction triggers const CompactThreshold = 0.80 -// AskFile is the main interaction file - write a prompt, read the response +// AskFile is the main interaction file - write a prompt, read the response. +// It implements FidAwareFile to provide per-fid session isolation. type AskFile struct { *protocol.BaseFile - client llm.Backend - mu sync.RWMutex - lastResponse string + sm *llm.SessionManager } // NewAskFile creates the ask file -func NewAskFile(client llm.Backend) *AskFile { +func NewAskFile(sm *llm.SessionManager) *AskFile { return &AskFile{ BaseFile: protocol.NewBaseFile("ask", 0666), - client: client, + sm: sm, } } +// Read implements File.Read (fallback for non-fid-aware access) func (f *AskFile) Read(p []byte, offset int64) (int, error) { - f.mu.RLock() - content := f.lastResponse - f.mu.RUnlock() + // Without fid context, we can't return session-specific data + // Return empty to indicate no data available + return 0, io.EOF +} + +// Write implements File.Write (fallback for non-fid-aware access) +func (f *AskFile) Write(p []byte, offset int64) (int, error) { + // Without fid context, we can't process the request properly + return 0, protocol.ErrPermission +} + +// ReadFid implements FidAwareFile.ReadFid +func (f *AskFile) ReadFid(fid uint32, p []byte, offset int64) (int, error) { + session := f.sm.GetOrCreate(fid) + content := session.LastResponse() // Add newline if not present if content != "" && !strings.HasSuffix(content, "\n") { @@ -47,7 +58,8 @@ func (f *AskFile) Read(p []byte, offset int64) (int, error) { return n, nil } -func (f *AskFile) Write(p []byte, offset int64) (int, error) { +// WriteFid implements FidAwareFile.WriteFid +func (f *AskFile) WriteFid(fid uint32, p []byte, offset int64) (int, error) { prompt := strings.TrimSpace(string(p)) if prompt == "" { return len(p), nil // Empty write is a no-op @@ -56,47 +68,39 @@ func (f *AskFile) Write(p []byte, offset int64) (int, error) { ctx := context.Background() // Check if we need to auto-compact before processing - tokens := f.client.TotalTokens() - limit := f.client.ContextLimit() + session := f.sm.GetOrCreate(fid) + tokens := session.TotalTokens() + limit := f.sm.ContextLimit() threshold := int(float64(limit) * CompactThreshold) if tokens > threshold { - log.Printf("llm9p: auto-compacting at %d/%d tokens (%.0f%% threshold)", - tokens, limit, CompactThreshold*100) - if err := f.client.Compact(ctx); err != nil { - log.Printf("llm9p: auto-compact failed: %v", err) - // Continue anyway - better to try than to fail - } else { - log.Printf("llm9p: auto-compact complete, now at %d tokens", - f.client.TotalTokens()) - } + log.Printf("llm9p: fid %d at %d/%d tokens (%.0f%% threshold) - consider resetting", + fid, tokens, limit, CompactThreshold*100) + // Note: Per-session compaction would require a different approach + // For now, just log a warning. Session reset via /new is the solution. } - // Make the API call - response, err := f.client.Ask(ctx, prompt) + // Make the API call using the session + _, err := f.sm.Ask(ctx, fid, prompt) if err != nil { - // Store error as response so it can be read - f.mu.Lock() - f.lastResponse = "Error: " + err.Error() - f.mu.Unlock() + // Error is stored in session.LastResponse by SessionManager return len(p), nil // Return success so client knows write completed } - f.mu.Lock() - f.lastResponse = response - f.mu.Unlock() - return len(p), nil } +// CloseFid implements FidAwareFile.CloseFid +func (f *AskFile) CloseFid(fid uint32) error { + // Clean up the session when the fid is clunked + f.sm.Remove(fid) + return nil +} + +// Stat returns the file's metadata func (f *AskFile) Stat() protocol.Stat { s := f.BaseFile.Stat() - f.mu.RLock() - content := f.lastResponse - f.mu.RUnlock() - if content != "" && !strings.HasSuffix(content, "\n") { - content += "\n" - } - s.Length = uint64(len(content)) + // Length is dynamic based on session, but without fid context we return 0 + s.Length = 0 return s } diff --git a/internal/llmfs/ask_test.go b/internal/llmfs/ask_test.go index 2df07986ed4f64621969a6c948e7aac1eaf61ad0..d179d1b9ae4b1321344cbd2e1aa90634cbf729a7 100644 --- a/internal/llmfs/ask_test.go +++ b/internal/llmfs/ask_test.go @@ -3,20 +3,23 @@ package llmfs import ( "io" "testing" + + "github.com/NERVsystems/llm9p/internal/llm" ) func TestAskFile_Read_Empty(t *testing.T) { mock := NewMockBackend() - ask := NewAskFile(mock) + sm := llm.NewSessionManager(mock) + ask := NewAskFile(sm) - // Initial read should be empty (EOF) + // Initial read with fid should be empty (EOF) buf := make([]byte, 100) - n, err := ask.Read(buf, 0) + n, err := ask.ReadFid(1, buf, 0) if err != io.EOF { - t.Errorf("Read() error = %v, want io.EOF", err) + t.Errorf("ReadFid() error = %v, want io.EOF", err) } if n != 0 { - t.Errorf("Read() n = %d, want 0", n) + t.Errorf("ReadFid() n = %d, want 0", n) } } @@ -24,49 +27,55 @@ func TestAskFile_WriteRead(t *testing.T) { mock := NewMockBackend() mock.askResponse = "Hello, I'm Claude!" - ask := NewAskFile(mock) + sm := llm.NewSessionManager(mock) + ask := NewAskFile(sm) + + fid := uint32(1) // Write a prompt prompt := "Hello!" - n, err := ask.Write([]byte(prompt), 0) + n, err := ask.WriteFid(fid, []byte(prompt), 0) if err != nil { - t.Fatalf("Write() error: %v", err) + t.Fatalf("WriteFid() error: %v", err) } if n != len(prompt) { - t.Errorf("Write() n = %d, want %d", n, len(prompt)) + t.Errorf("WriteFid() n = %d, want %d", n, len(prompt)) } // Read response buf := make([]byte, 100) - readN, err := ask.Read(buf, 0) + readN, err := ask.ReadFid(fid, buf, 0) if err != nil { - t.Fatalf("Read() error: %v", err) + t.Fatalf("ReadFid() error: %v", err) } response := string(buf[:readN]) expected := "Hello, I'm Claude!\n" if response != expected { - t.Errorf("Read() = %q, want %q", response, expected) + t.Errorf("ReadFid() = %q, want %q", response, expected) } } func TestAskFile_Write_EmptyNoOp(t *testing.T) { mock := NewMockBackend() - ask := NewAskFile(mock) + sm := llm.NewSessionManager(mock) + ask := NewAskFile(sm) + + fid := uint32(1) // Empty write should be no-op - n, err := ask.Write([]byte(""), 0) + n, err := ask.WriteFid(fid, []byte(""), 0) if err != nil { - t.Fatalf("Write() error: %v", err) + t.Fatalf("WriteFid() error: %v", err) } if n != 0 { - t.Errorf("Write('') n = %d, want 0", n) + t.Errorf("WriteFid('') n = %d, want 0", n) } // Whitespace-only also no-op - n, err = ask.Write([]byte(" \n\t "), 0) + n, err = ask.WriteFid(fid, []byte(" \n\t "), 0) if err != nil { - t.Fatalf("Write() error: %v", err) + t.Fatalf("WriteFid() error: %v", err) } } @@ -74,88 +83,89 @@ func TestAskFile_Write_Error(t *testing.T) { mock := NewMockBackend() mock.askError = io.ErrUnexpectedEOF - ask := NewAskFile(mock) + sm := llm.NewSessionManager(mock) + ask := NewAskFile(sm) + + fid := uint32(1) // Write should succeed (error is stored for read) - n, err := ask.Write([]byte("test"), 0) + n, err := ask.WriteFid(fid, []byte("test"), 0) if err != nil { - t.Fatalf("Write() error: %v", err) + t.Fatalf("WriteFid() error: %v", err) } if n != 4 { - t.Errorf("Write() n = %d, want 4", n) + t.Errorf("WriteFid() n = %d, want 4", n) } // Read should return the error buf := make([]byte, 100) - readN, _ := ask.Read(buf, 0) + readN, _ := ask.ReadFid(fid, buf, 0) response := string(buf[:readN]) - if response[:6] != "Error:" { - t.Errorf("Read() = %q, should start with 'Error:'", response) + if len(response) < 6 || response[:6] != "Error:" { + t.Errorf("ReadFid() = %q, should start with 'Error:'", response) } } -func TestAskFile_AutoCompaction_BelowThreshold(t *testing.T) { +func TestAskFile_SessionIsolation(t *testing.T) { mock := NewMockBackend() - mock.totalTokens = 100000 // 50% of 200K, below 80% threshold - mock.contextLimit = 200000 mock.askResponse = "response" - ask := NewAskFile(mock) + sm := llm.NewSessionManager(mock) + ask := NewAskFile(sm) - // Write should NOT trigger compaction - ask.Write([]byte("test"), 0) + fid1 := uint32(1) + fid2 := uint32(2) - if mock.compactCalled { - t.Error("Compact() should not be called below threshold") - } -} + // Write to fid1 + mock.askResponse = "response for fid1" + ask.WriteFid(fid1, []byte("prompt1"), 0) -func TestAskFile_AutoCompaction_AboveThreshold(t *testing.T) { - mock := NewMockBackend() - mock.totalTokens = 170000 // 85% of 200K, above 80% threshold - mock.contextLimit = 200000 - mock.askResponse = "response" - - ask := NewAskFile(mock) + // Write to fid2 + mock.askResponse = "response for fid2" + ask.WriteFid(fid2, []byte("prompt2"), 0) - // Write SHOULD trigger compaction - ask.Write([]byte("test"), 0) + // Read from fid1 - should get fid1's response + buf := make([]byte, 100) + n, _ := ask.ReadFid(fid1, buf, 0) + response1 := string(buf[:n]) + if response1 != "response for fid1\n" { + t.Errorf("fid1 ReadFid() = %q, want %q", response1, "response for fid1\n") + } - if !mock.compactCalled { - t.Error("Compact() should be called when above threshold") + // Read from fid2 - should get fid2's response + n, _ = ask.ReadFid(fid2, buf, 0) + response2 := string(buf[:n]) + if response2 != "response for fid2\n" { + t.Errorf("fid2 ReadFid() = %q, want %q", response2, "response for fid2\n") } } -func TestAskFile_AutoCompaction_ExactThreshold(t *testing.T) { +func TestAskFile_CloseFid(t *testing.T) { mock := NewMockBackend() - mock.totalTokens = 160000 // Exactly 80% of 200K - mock.contextLimit = 200000 mock.askResponse = "response" - ask := NewAskFile(mock) + sm := llm.NewSessionManager(mock) + ask := NewAskFile(sm) - // Write should NOT trigger (we use > not >=) - ask.Write([]byte("test"), 0) + fid := uint32(1) - if mock.compactCalled { - t.Error("Compact() should not be called at exactly threshold") - } -} - -func TestAskFile_AutoCompaction_JustAboveThreshold(t *testing.T) { - mock := NewMockBackend() - mock.totalTokens = 160001 // Just over 80% - mock.contextLimit = 200000 - mock.askResponse = "response" + // Write to create session + ask.WriteFid(fid, []byte("test"), 0) - ask := NewAskFile(mock) + // Session should exist + session := sm.Get(fid) + if session == nil { + t.Fatal("session should exist after WriteFid") + } - // Write SHOULD trigger compaction - ask.Write([]byte("test"), 0) + // Close the fid + ask.CloseFid(fid) - if !mock.compactCalled { - t.Error("Compact() should be called when just above threshold") + // Session should be removed + session = sm.Get(fid) + if session != nil { + t.Error("session should be removed after CloseFid") } } @@ -163,35 +173,29 @@ func TestAskFile_Stat(t *testing.T) { mock := NewMockBackend() mock.askResponse = "Hello!" - ask := NewAskFile(mock) + sm := llm.NewSessionManager(mock) + ask := NewAskFile(sm) - // Initial stat - no response yet + // Stat without fid context returns 0 length stat := ask.Stat() if stat.Length != 0 { - t.Errorf("Stat().Length = %d, want 0 (no response yet)", stat.Length) - } - - // Write to get a response - ask.Write([]byte("test"), 0) - - // Stat should now show response length (with newline) - stat = ask.Stat() - expected := uint64(len("Hello!\n")) - if stat.Length != expected { - t.Errorf("Stat().Length = %d, want %d", stat.Length, expected) + t.Errorf("Stat().Length = %d, want 0", stat.Length) } } func TestAskFile_ResponseNewline(t *testing.T) { mock := NewMockBackend() + sm := llm.NewSessionManager(mock) + + fid := uint32(1) // Response without trailing newline mock.askResponse = "No newline" - ask := NewAskFile(mock) - ask.Write([]byte("test"), 0) + ask := NewAskFile(sm) + ask.WriteFid(fid, []byte("test"), 0) buf := make([]byte, 100) - n, _ := ask.Read(buf, 0) + n, _ := ask.ReadFid(fid, buf, 0) response := string(buf[:n]) // Should have newline added @@ -199,12 +203,14 @@ func TestAskFile_ResponseNewline(t *testing.T) { t.Errorf("Response should end with newline, got %q", response) } + // New session for second test + fid2 := uint32(2) + // Response with trailing newline already mock.askResponse = "Has newline\n" - ask2 := NewAskFile(mock) - ask2.Write([]byte("test"), 0) + ask.WriteFid(fid2, []byte("test"), 0) - n, _ = ask2.Read(buf, 0) + n, _ = ask.ReadFid(fid2, buf, 0) response = string(buf[:n]) // Should NOT double the newline diff --git a/internal/llmfs/context.go b/internal/llmfs/context.go index d7b3bf2b5ec3bf89d188927b63ad50dae63b0eed..32ed57a9b9356323acde975b64fdca2088e494c9 100644 --- a/internal/llmfs/context.go +++ b/internal/llmfs/context.go @@ -8,24 +8,39 @@ import ( "github.com/NERVsystems/llm9p/internal/protocol" ) -// ContextFile exposes the conversation history -// Read: returns JSON of conversation history -// Write: appends a system message to context +// ContextFile exposes the conversation history. +// It implements FidAwareFile to provide per-fid session isolation. +// Read: returns JSON of conversation history for this fid's session +// Write: appends a system message to this fid's session context type ContextFile struct { *protocol.BaseFile - client llm.Backend + sm *llm.SessionManager } // NewContextFile creates the context file -func NewContextFile(client llm.Backend) *ContextFile { +func NewContextFile(sm *llm.SessionManager) *ContextFile { return &ContextFile{ BaseFile: protocol.NewBaseFile("context", 0666), - client: client, + sm: sm, } } +// Read implements File.Read (fallback for non-fid-aware access) func (f *ContextFile) Read(p []byte, offset int64) (int, error) { - content, err := f.client.MessagesJSON() + // Without fid context, return empty + return 0, io.EOF +} + +// Write implements File.Write (fallback for non-fid-aware access) +func (f *ContextFile) Write(p []byte, offset int64) (int, error) { + // Without fid context, we can't add to a specific session + return 0, protocol.ErrPermission +} + +// ReadFid implements FidAwareFile.ReadFid +func (f *ContextFile) ReadFid(fid uint32, p []byte, offset int64) (int, error) { + session := f.sm.GetOrCreate(fid) + content, err := session.MessagesJSON() if err != nil { return 0, err } @@ -39,18 +54,28 @@ func (f *ContextFile) Read(p []byte, offset int64) (int, error) { return n, nil } -func (f *ContextFile) Write(p []byte, offset int64) (int, error) { - // Writing appends a system message to the context +// WriteFid implements FidAwareFile.WriteFid +func (f *ContextFile) WriteFid(fid uint32, p []byte, offset int64) (int, error) { + // Writing appends a system message to this fid's session context msg := strings.TrimSpace(string(p)) if msg != "" { - f.client.AddSystemMessage(msg) + session := f.sm.GetOrCreate(fid) + session.AddSystemMessage(msg) } return len(p), nil } +// CloseFid implements FidAwareFile.CloseFid +func (f *ContextFile) CloseFid(fid uint32) error { + // No per-fid state to clean up for this file + // (session cleanup is handled by AskFile) + return nil +} + +// Stat returns the file's metadata func (f *ContextFile) Stat() protocol.Stat { s := f.BaseFile.Stat() - content, _ := f.client.MessagesJSON() - s.Length = uint64(len(content) + 1) // +1 for newline + // Length is dynamic based on session, but without fid context we return 0 + s.Length = 0 return s } diff --git a/internal/llmfs/mock_backend_test.go b/internal/llmfs/mock_backend_test.go index a8b198db98a63e1f03ea257f77dd9643fbc42104..e9a5d6ae76efbc7b9f1301d9bfdb9a8834e04018 100644 --- a/internal/llmfs/mock_backend_test.go +++ b/internal/llmfs/mock_backend_test.go @@ -12,6 +12,7 @@ type MockBackend struct { model string temperature float64 systemPrompt string + prefill string messages []llm.Message lastTokens int totalTokens int @@ -46,6 +47,8 @@ func (m *MockBackend) SystemPrompt() string { return m.systemPrompt } func (m *MockBackend) SetSystemPrompt(prompt string) { m.systemPrompt = prompt } func (m *MockBackend) ThinkingTokens() int { return m.thinkingTokens } func (m *MockBackend) SetThinkingTokens(tokens int) { m.thinkingTokens = tokens } +func (m *MockBackend) Prefill() string { return m.prefill } +func (m *MockBackend) SetPrefill(prefill string) { m.prefill = prefill } func (m *MockBackend) LastTokens() int { return m.lastTokens } func (m *MockBackend) TotalTokens() int { return m.totalTokens } func (m *MockBackend) ContextLimit() int { return m.contextLimit } @@ -92,6 +95,14 @@ func (m *MockBackend) Ask(ctx context.Context, prompt string) (string, error) { return m.askResponse, nil } +func (m *MockBackend) AskWithHistory(ctx context.Context, history []llm.Message, prompt string) (string, int, error) { + if m.askError != nil { + return "", 0, m.askError + } + tokens := len(prompt) + len(m.askResponse) + return m.askResponse, tokens, nil +} + func (m *MockBackend) StartStream(ctx context.Context, prompt string) error { return fmt.Errorf("streaming not implemented in mock") } diff --git a/internal/llmfs/new.go b/internal/llmfs/new.go index 4093880fd2430557b55d71e5c881b4ec6a8dc8f4..d11cf5c03e720f0f9797e3de0ee3e06587b51ca0 100644 --- a/internal/llmfs/new.go +++ b/internal/llmfs/new.go @@ -5,30 +5,51 @@ import ( "github.com/NERVsystems/llm9p/internal/protocol" ) -// NewFile is a write-only file that resets the conversation when written to +// NewFile is a write-only file that resets the conversation when written to. +// It implements FidAwareFile to reset only the session for the writing fid. type NewFile struct { *protocol.BaseFile - client llm.Backend + sm *llm.SessionManager } // NewNewFile creates the new file -func NewNewFile(client llm.Backend) *NewFile { +func NewNewFile(sm *llm.SessionManager) *NewFile { return &NewFile{ BaseFile: protocol.NewBaseFile("new", 0222), - client: client, + sm: sm, } } +// Read implements File.Read func (f *NewFile) Read(p []byte, offset int64) (int, error) { return 0, protocol.ErrPermission } +// Write implements File.Write (fallback for non-fid-aware access) func (f *NewFile) Write(p []byte, offset int64) (int, error) { - // Any write resets the conversation - f.client.Reset() + // Without fid context, we can't reset a specific session + return 0, protocol.ErrPermission +} + +// ReadFid implements FidAwareFile.ReadFid +func (f *NewFile) ReadFid(fid uint32, p []byte, offset int64) (int, error) { + return 0, protocol.ErrPermission +} + +// WriteFid implements FidAwareFile.WriteFid +func (f *NewFile) WriteFid(fid uint32, p []byte, offset int64) (int, error) { + // Reset only the session for this fid + f.sm.Reset(fid) return len(p), nil } +// CloseFid implements FidAwareFile.CloseFid +func (f *NewFile) CloseFid(fid uint32) error { + // No per-fid state to clean up for this file + return nil +} + +// Stat returns the file's metadata func (f *NewFile) Stat() protocol.Stat { s := f.BaseFile.Stat() s.Length = 0 diff --git a/internal/llmfs/prefill.go b/internal/llmfs/prefill.go new file mode 100644 index 0000000000000000000000000000000000000000..4ef3cd25304e6da253ad893de87d3dc703db6fc0 --- /dev/null +++ b/internal/llmfs/prefill.go @@ -0,0 +1,54 @@ +package llmfs + +import ( + "io" + "strings" + + "github.com/NERVsystems/llm9p/internal/llm" + "github.com/NERVsystems/llm9p/internal/protocol" +) + +// PrefillFile exposes the assistant response prefill (read/write). +// Prefill helps keep the model in character by prepending a string +// to the assistant's response (e.g., "[Veltro] "). +type PrefillFile struct { + *protocol.BaseFile + client llm.Backend +} + +// NewPrefillFile creates the prefill file +func NewPrefillFile(client llm.Backend) *PrefillFile { + return &PrefillFile{ + BaseFile: protocol.NewBaseFile("prefill", 0666), + client: client, + } +} + +func (f *PrefillFile) Read(p []byte, offset int64) (int, error) { + content := f.client.Prefill() + if content != "" { + content += "\n" + } + if offset >= int64(len(content)) { + return 0, io.EOF + } + n := copy(p, content[offset:]) + return n, nil +} + +func (f *PrefillFile) Write(p []byte, offset int64) (int, error) { + prefill := strings.TrimSpace(string(p)) + f.client.SetPrefill(prefill) + return len(p), nil +} + +func (f *PrefillFile) Stat() protocol.Stat { + s := f.BaseFile.Stat() + content := f.client.Prefill() + if content != "" { + s.Length = uint64(len(content) + 1) // +1 for newline + } else { + s.Length = 0 + } + return s +} diff --git a/internal/llmfs/root.go b/internal/llmfs/root.go index dbbc1f13cd2309a57ddb92571ac4efef2ac7fe90..066676e8c3b7d395358fcf40a84b16372110074d 100644 --- a/internal/llmfs/root.go +++ b/internal/llmfs/root.go @@ -6,27 +6,38 @@ import ( "github.com/NERVsystems/llm9p/internal/protocol" ) -// NewRoot creates the root directory of the LLM filesystem -func NewRoot(client llm.Backend) protocol.Dir { +// NewRoot creates the root directory of the LLM filesystem. +// It takes a SessionManager which provides per-fid session isolation +// and access to the underlying backend for global settings. +func NewRoot(sm *llm.SessionManager) protocol.Dir { + backend := sm.Backend() root := protocol.NewStaticDir("llm") - // Add all files - root.AddChild(NewAskFile(client)) - root.AddChild(NewModelFile(client)) - root.AddChild(NewTemperatureFile(client)) - root.AddChild(NewSystemFile(client)) - root.AddChild(NewTokensFile(client)) - root.AddChild(NewNewFile(client)) - root.AddChild(NewContextFile(client)) - root.AddChild(NewThinkingFile(client)) + // Session-aware files (per-fid isolation) + root.AddChild(NewAskFile(sm)) + root.AddChild(NewNewFile(sm)) + root.AddChild(NewContextFile(sm)) + + // Global settings files (shared across all fids) + root.AddChild(NewModelFile(backend)) + root.AddChild(NewTemperatureFile(backend)) + root.AddChild(NewSystemFile(backend)) + root.AddChild(NewThinkingFile(backend)) + root.AddChild(NewPrefillFile(backend)) + + // Token tracking (uses backend's global counters) + root.AddChild(NewTokensFile(backend)) + root.AddChild(NewUsageFile(backend)) + root.AddChild(NewMetricsFile(backend)) + root.AddChild(NewCompactFile(backend)) + + // Static files root.AddChild(NewExampleFile()) - root.AddChild(NewUsageFile(client)) - root.AddChild(NewCompactFile(client)) - // Add stream directory + // Add stream directory (uses backend directly) streamDir := protocol.NewStaticDir("stream") - streamDir.AddChild(NewStreamAskFile(client)) - streamDir.AddChild(NewChunkFile(client)) + streamDir.AddChild(NewStreamAskFile(backend)) + streamDir.AddChild(NewChunkFile(backend)) root.AddChild(streamDir) return root diff --git a/internal/protocol/fs.go b/internal/protocol/fs.go index e137af1e1b5f9feb873a491228890f42138b4494..3d7608a916a2d954afc3c725f6665728012fda47 100644 --- a/internal/protocol/fs.go +++ b/internal/protocol/fs.go @@ -36,6 +36,22 @@ type Dir interface { Lookup(name string) (File, error) } +// FidAwareFile is implemented by files that need per-fid state. +// When a file implements this interface, the server will call the +// fid-aware methods instead of the standard File methods. +type FidAwareFile interface { + File + + // ReadFid reads from the file with fid context + ReadFid(fid uint32, p []byte, offset int64) (n int, err error) + + // WriteFid writes to the file with fid context + WriteFid(fid uint32, p []byte, offset int64) (n int, err error) + + // CloseFid is called when a fid is clunked + CloseFid(fid uint32) error +} + // pathCounter generates unique path IDs for qids var pathCounter uint64 diff --git a/internal/protocol/server.go b/internal/protocol/server.go index 123c52fb9f692018cd9b490818907379618667db..5eb0f6e25a2017e9c8558352f62b8ce7dc0fcf39 100644 --- a/internal/protocol/server.go +++ b/internal/protocol/server.go @@ -272,7 +272,15 @@ func (s *Server) handleRead(state *clientState, payload []byte, buf []byte) ([]b } data := make([]byte, count) - n, err := file.Read(data, int64(msg.Offset)) + var n int + + // Check for fid-aware file + if faf, ok := file.(FidAwareFile); ok { + n, err = faf.ReadFid(msg.Fid, data, int64(msg.Offset)) + } else { + n, err = file.Read(data, int64(msg.Offset)) + } + if err != nil && err != io.EOF { return s.errorResponse(buf, err.Error()) } @@ -293,7 +301,15 @@ func (s *Server) handleWrite(state *clientState, payload []byte, buf []byte) ([] return s.errorResponse(buf, ErrBadFid.Error()) } - n, err := file.Write(msg.Data, int64(msg.Offset)) + var n int + + // Check for fid-aware file + if faf, ok := file.(FidAwareFile); ok { + n, err = faf.WriteFid(msg.Fid, msg.Data, int64(msg.Offset)) + } else { + n, err = file.Write(msg.Data, int64(msg.Offset)) + } + if err != nil { return s.errorResponse(buf, err.Error()) } @@ -314,6 +330,11 @@ func (s *Server) handleClunk(state *clientState, payload []byte, buf []byte) ([] return s.errorResponse(buf, ErrBadFid.Error()) } + // Call CloseFid for fid-aware files to clean up per-fid state + if faf, ok := file.(FidAwareFile); ok { + faf.CloseFid(msg.Fid) + } + file.Close() delete(state.fids, msg.Fid)