diff --git a/protocol/messages.go b/protocol/messages.go index 44f233a..728f7c9 100644 --- a/protocol/messages.go +++ b/protocol/messages.go @@ -108,9 +108,11 @@ const ( TaskOpGrep TaskOperation = "grep" TaskOpBash TaskOperation = "bash" TaskOpWebFetch TaskOperation = "webfetch" - TaskOpMCPCall TaskOperation = "mcp_call" - TaskOpMCPListTools TaskOperation = "mcp_list_tools" - TaskOpSyncSkill TaskOperation = "sync_skill" + TaskOpMCPCall TaskOperation = "mcp_call" + TaskOpMCPListTools TaskOperation = "mcp_list_tools" + TaskOpSyncSkill TaskOperation = "sync_skill" + TaskOpStageKnowledgeFiles TaskOperation = "stage_knowledge_files" + TaskOpDeleteKnowledgeFiles TaskOperation = "delete_knowledge_files" ) // TaskRequestPayload is the payload for task request messages. @@ -323,3 +325,37 @@ type MCPResultPayload struct { Result json.RawMessage `json:"result,omitempty"` Error string `json:"error,omitempty"` } + +// KnowledgeFile is a single file entry in a stage_knowledge_files request. +type KnowledgeFile struct { + RelPath string `json:"rel_path"` + Checksum string `json:"checksum"` + ContentB64 string `json:"content_b64"` +} + +// StageKnowledgeFilesArgs are the arguments for stage_knowledge_files operation. +type StageKnowledgeFilesArgs struct { + Files []KnowledgeFile `json:"files"` +} + +// KnowledgeFileStatus is the per-file result entry in the stage ack. +type KnowledgeFileStatus struct { + RelPath string `json:"rel_path"` + Success bool `json:"success"` + Error string `json:"error,omitempty"` +} + +// StageKnowledgeFilesResult is the result of a stage_knowledge_files operation. +type StageKnowledgeFilesResult struct { + Files []KnowledgeFileStatus `json:"files"` +} + +// DeleteKnowledgeFilesArgs are the arguments for delete_knowledge_files operation. +type DeleteKnowledgeFilesArgs struct { + RelPaths []string `json:"rel_paths"` +} + +// DeleteKnowledgeFilesResult is the result of a delete_knowledge_files operation. +type DeleteKnowledgeFilesResult struct { + Success bool `json:"success"` +} diff --git a/workspace/knowledge.go b/workspace/knowledge.go new file mode 100644 index 0000000..9cc4a2e --- /dev/null +++ b/workspace/knowledge.go @@ -0,0 +1,249 @@ +package workspace + +import ( + "context" + "encoding/base64" + "encoding/json" + "fmt" + "log/slog" + "os" + "path/filepath" + "strings" + "syscall" + + "github.com/flashcatcloud/flashduty-runner/protocol" +) + +const ( + // sentinelName is the hidden JSON map that tracks staged-file checksums. + // Safari reads this to decide which knowledge pack files are already current. + sentinelName = ".safari-knowledge-sentinel.json" +) + +// validateKnowledgeRelPath enforces the path rules for knowledge file operations. +// +// Rules (from the Safari-side contract): +// - Must not contain path separators or double-dot components — the runner +// only writes flat files in the workspace root, never in sub-directories. +// - Leading-dot filenames are rejected because they are hidden by convention; +// the sentinel is written by the runner itself and is never staged by clients. +func validateKnowledgeRelPath(relPath string) error { + if relPath == "" { + return fmt.Errorf("rel_path must not be empty") + } + if strings.ContainsAny(relPath, `/\`) { + return fmt.Errorf("rel_path must not contain path separators: %q", relPath) + } + // Reject the bare ".." token. Slash-separated traversal like "foo/../bar" + // is already blocked above, but a plain ".." with no slashes still escapes. + if relPath == ".." { + return fmt.Errorf("rel_path must not be '..': %q", relPath) + } + if strings.HasPrefix(relPath, ".") { + // Hidden files (including the sentinel itself) cannot be staged by clients. + // The runner owns the sentinel exclusively. + return fmt.Errorf("rel_path must not start with '.': %q", relPath) + } + return nil +} + +// atomicWriteFile writes data to path using a temp-file + rename so readers +// never see a partially-written file. The temp file is created in the same +// directory as the target to guarantee the rename stays on the same filesystem. +func atomicWriteFile(path string, data []byte, perm os.FileMode) error { + dir := filepath.Dir(path) + tmp, err := os.CreateTemp(dir, ".tmp-knowledge-*") + if err != nil { + return fmt.Errorf("failed to create temp file: %w", err) + } + tmpName := tmp.Name() + + // Clean up the temp file on any error path. + var writeErr error + defer func() { + if writeErr != nil { + _ = os.Remove(tmpName) + } + }() + + if _, writeErr = tmp.Write(data); writeErr != nil { + _ = tmp.Close() + return fmt.Errorf("failed to write temp file: %w", writeErr) + } + if writeErr = tmp.Chmod(perm); writeErr != nil { + _ = tmp.Close() + return fmt.Errorf("failed to chmod temp file: %w", writeErr) + } + if writeErr = tmp.Close(); writeErr != nil { + return fmt.Errorf("failed to close temp file: %w", writeErr) + } + + if writeErr = os.Rename(tmpName, path); writeErr != nil { + return fmt.Errorf("failed to rename temp file: %w", writeErr) + } + return nil +} + +// withSentinelLock opens (or creates) the sentinel file, acquires an exclusive +// advisory flock on it, calls fn, then releases the lock. The advisory lock +// protects concurrent read-modify-write cycles across BYOC sessions that share +// the same filesystem (e.g. multiple Safari instances writing to the same +// worknode workspace root). +// +// Note: syscall.Flock is available on Linux and macOS. The runner is deployed +// on those platforms only; Windows is not supported today. +func withSentinelLock(sentinelPath string, fn func() error) error { + // Open or create the sentinel file just to acquire the lock fd. + // We do NOT read/write through this fd to keep flock + atomic-write + // concerns separate. + lockFile, err := os.OpenFile(sentinelPath, os.O_RDWR|os.O_CREATE, 0o600) + if err != nil { + return fmt.Errorf("failed to open sentinel for locking: %w", err) + } + defer func() { + _ = lockFile.Close() + }() + + fd := int(lockFile.Fd()) //nolint:gosec // os.File.Fd returns uintptr but the underlying OS fd is always a valid int on unix + if err := syscall.Flock(fd, syscall.LOCK_EX); err != nil { + return fmt.Errorf("failed to acquire sentinel lock: %w", err) + } + defer func() { + _ = syscall.Flock(fd, syscall.LOCK_UN) + }() + + return fn() +} + +// readSentinel reads the sentinel JSON map. Missing file or empty file → +// empty map (both are expected states on first use). Corrupt JSON → log a +// warning, return empty map (safe rebuild on next stage). +func readSentinel(sentinelPath string) map[string]string { + data, err := os.ReadFile(sentinelPath) + if err != nil { + if !os.IsNotExist(err) { + slog.Warn("failed to read sentinel, treating as empty", "error", err) + } + return make(map[string]string) + } + // An empty file is the initial state created by withSentinelLock; treat it + // as an empty map without logging a warning. + if len(data) == 0 { + return make(map[string]string) + } + var m map[string]string + if err := json.Unmarshal(data, &m); err != nil { + slog.Warn("sentinel JSON corrupt, treating as empty", "error", err) + return make(map[string]string) + } + return m +} + +// writeSentinel atomically rewrites the sentinel with the given map. +func writeSentinel(sentinelPath string, m map[string]string) error { + data, err := json.Marshal(m) + if err != nil { + return fmt.Errorf("failed to marshal sentinel: %w", err) + } + return atomicWriteFile(sentinelPath, data, 0o644) +} + +// StageKnowledgeFiles writes the supplied files into the workspace root and +// updates the sentinel checksum map. +func (w *Workspace) StageKnowledgeFiles(ctx context.Context, args *protocol.StageKnowledgeFilesArgs) (*protocol.StageKnowledgeFilesResult, error) { + result := &protocol.StageKnowledgeFilesResult{ + Files: make([]protocol.KnowledgeFileStatus, 0, len(args.Files)), + } + + // validated collects (relPath, checksum) for files that landed successfully; + // we only merge these into the sentinel. + type staged struct{ relPath, checksum string } + var succeeded []staged + + for _, f := range args.Files { + status := protocol.KnowledgeFileStatus{RelPath: f.RelPath} + + if err := validateKnowledgeRelPath(f.RelPath); err != nil { + status.Success = false + status.Error = err.Error() + result.Files = append(result.Files, status) + continue + } + + content, err := base64.StdEncoding.DecodeString(f.ContentB64) + if err != nil { + status.Success = false + status.Error = fmt.Sprintf("failed to decode content_b64: %v", err) + result.Files = append(result.Files, status) + continue + } + + targetPath := filepath.Join(w.root, f.RelPath) + if err := atomicWriteFile(targetPath, content, 0o644); err != nil { + status.Success = false + status.Error = err.Error() + result.Files = append(result.Files, status) + continue + } + + status.Success = true + result.Files = append(result.Files, status) + succeeded = append(succeeded, staged{f.RelPath, f.Checksum}) + } + + // Update sentinel for successfully written files under advisory lock. + if len(succeeded) > 0 { + sentinelPath := filepath.Join(w.root, sentinelName) + if err := withSentinelLock(sentinelPath, func() error { + m := readSentinel(sentinelPath) + for _, s := range succeeded { + m[s.relPath] = s.checksum + } + return writeSentinel(sentinelPath, m) + }); err != nil { + // Sentinel write failure is non-fatal for the already-written files, + // but we log it clearly so operators can investigate. + slog.Error("failed to update sentinel after staging", "error", err) + } + } + + return result, nil +} + +// DeleteKnowledgeFiles removes the supplied files from the workspace root and +// scrubs their entries from the sentinel. +func (w *Workspace) DeleteKnowledgeFiles(ctx context.Context, args *protocol.DeleteKnowledgeFilesArgs) (*protocol.DeleteKnowledgeFilesResult, error) { + // toRemove collects validated rel_paths so we can clean the sentinel in one + // locked pass after the file removals. + toRemove := make(map[string]struct{}, len(args.RelPaths)) + + for _, relPath := range args.RelPaths { + if err := validateKnowledgeRelPath(relPath); err != nil { + slog.Warn("skipping invalid rel_path in delete_knowledge_files", "rel_path", relPath, "error", err) + continue + } + toRemove[relPath] = struct{}{} + + targetPath := filepath.Join(w.root, relPath) + if err := os.Remove(targetPath); err != nil && !os.IsNotExist(err) { + slog.Warn("failed to remove knowledge file", "rel_path", relPath, "error", err) + } + } + + // Remove entries from sentinel for all valid paths (idempotent — missing + // sentinel entries are simply no-ops in the delete loop). + if len(toRemove) > 0 { + sentinelPath := filepath.Join(w.root, sentinelName) + if err := withSentinelLock(sentinelPath, func() error { + m := readSentinel(sentinelPath) + for rp := range toRemove { + delete(m, rp) + } + return writeSentinel(sentinelPath, m) + }); err != nil { + slog.Error("failed to update sentinel after deletion", "error", err) + } + } + + return &protocol.DeleteKnowledgeFilesResult{Success: true}, nil +} diff --git a/workspace/knowledge_test.go b/workspace/knowledge_test.go new file mode 100644 index 0000000..8d4c811 --- /dev/null +++ b/workspace/knowledge_test.go @@ -0,0 +1,238 @@ +package workspace + +import ( + "context" + "encoding/base64" + "encoding/json" + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/flashcatcloud/flashduty-runner/protocol" +) + +// b64 encodes a string to base64 — reduces noise in test table declarations. +func b64(s string) string { return base64.StdEncoding.EncodeToString([]byte(s)) } + +// readSentinelMap reads the sentinel JSON file from root and returns the map. +func readSentinelMap(t *testing.T, root string) map[string]string { + t.Helper() + data, err := os.ReadFile(filepath.Join(root, sentinelName)) + require.NoError(t, err) + var m map[string]string + require.NoError(t, json.Unmarshal(data, &m)) + return m +} + +// fileContent reads a file and returns its content as a string. +func fileContent(t *testing.T, root, relPath string) string { + t.Helper() + data, err := os.ReadFile(filepath.Join(root, relPath)) + require.NoError(t, err) + return string(data) +} + +// Case 1: stage one file → file lands, sentinel has the entry. +func TestStageKnowledgeFiles_Single(t *testing.T) { + ws := newTestWorkspace(t) + ctx := context.Background() + + result, err := ws.StageKnowledgeFiles(ctx, &protocol.StageKnowledgeFilesArgs{ + Files: []protocol.KnowledgeFile{ + {RelPath: "DUTY.md", Checksum: "abc123", ContentB64: b64("# Duty\n")}, + }, + }) + require.NoError(t, err) + require.Len(t, result.Files, 1) + assert.True(t, result.Files[0].Success) + assert.Empty(t, result.Files[0].Error) + + assert.Equal(t, "# Duty\n", fileContent(t, ws.Root(), "DUTY.md")) + + sentinel := readSentinelMap(t, ws.Root()) + assert.Equal(t, "abc123", sentinel["DUTY.md"]) +} + +// Case 2: mixed batch — one valid file + one with '/' in rel_path. +// The valid file must land; the invalid one must return an error in the ack; +// the sentinel must contain only the valid entry. +func TestStageKnowledgeFiles_MixedValidity(t *testing.T) { + ws := newTestWorkspace(t) + ctx := context.Background() + + result, err := ws.StageKnowledgeFiles(ctx, &protocol.StageKnowledgeFilesArgs{ + Files: []protocol.KnowledgeFile{ + {RelPath: "DUTY.md", Checksum: "goodsum", ContentB64: b64("content")}, + {RelPath: "sub/evil.md", Checksum: "badsum", ContentB64: b64("evil")}, + }, + }) + require.NoError(t, err) + require.Len(t, result.Files, 2) + + // First entry (valid) should succeed. + assert.True(t, result.Files[0].Success) + assert.Empty(t, result.Files[0].Error) + + // Second entry (slash in path) should fail. + assert.False(t, result.Files[1].Success) + assert.NotEmpty(t, result.Files[1].Error) + + // Only valid file exists on disk. + assert.FileExists(t, filepath.Join(ws.Root(), "DUTY.md")) + assert.NoFileExists(t, filepath.Join(ws.Root(), "sub/evil.md")) + + // Sentinel has only the good entry. + sentinel := readSentinelMap(t, ws.Root()) + assert.Equal(t, "goodsum", sentinel["DUTY.md"]) + _, hasBad := sentinel["sub/evil.md"] + assert.False(t, hasBad) +} + +// Case 3: stage the same file twice with different checksums → second stage +// overwrites content and updates the sentinel checksum. +func TestStageKnowledgeFiles_Overwrite(t *testing.T) { + ws := newTestWorkspace(t) + ctx := context.Background() + + _, err := ws.StageKnowledgeFiles(ctx, &protocol.StageKnowledgeFilesArgs{ + Files: []protocol.KnowledgeFile{ + {RelPath: "runbook.md", Checksum: "v1sum", ContentB64: b64("version 1")}, + }, + }) + require.NoError(t, err) + + _, err = ws.StageKnowledgeFiles(ctx, &protocol.StageKnowledgeFilesArgs{ + Files: []protocol.KnowledgeFile{ + {RelPath: "runbook.md", Checksum: "v2sum", ContentB64: b64("version 2")}, + }, + }) + require.NoError(t, err) + + assert.Equal(t, "version 2", fileContent(t, ws.Root(), "runbook.md")) + + sentinel := readSentinelMap(t, ws.Root()) + assert.Equal(t, "v2sum", sentinel["runbook.md"]) +} + +// Case 4: delete a previously staged file → file gone, sentinel entry gone. +func TestDeleteKnowledgeFiles_Staged(t *testing.T) { + ws := newTestWorkspace(t) + ctx := context.Background() + + // Stage first. + _, err := ws.StageKnowledgeFiles(ctx, &protocol.StageKnowledgeFilesArgs{ + Files: []protocol.KnowledgeFile{ + {RelPath: "old-runbook.md", Checksum: "deadbeef", ContentB64: b64("old content")}, + }, + }) + require.NoError(t, err) + + // Delete. + delResult, err := ws.DeleteKnowledgeFiles(ctx, &protocol.DeleteKnowledgeFilesArgs{ + RelPaths: []string{"old-runbook.md"}, + }) + require.NoError(t, err) + assert.True(t, delResult.Success) + + assert.NoFileExists(t, filepath.Join(ws.Root(), "old-runbook.md")) + + sentinel := readSentinelMap(t, ws.Root()) + _, exists := sentinel["old-runbook.md"] + assert.False(t, exists) +} + +// Case 5: delete a path that was never staged → no error (idempotent). +func TestDeleteKnowledgeFiles_Missing(t *testing.T) { + ws := newTestWorkspace(t) + ctx := context.Background() + + result, err := ws.DeleteKnowledgeFiles(ctx, &protocol.DeleteKnowledgeFilesArgs{ + RelPaths: []string{"does-not-exist.md"}, + }) + require.NoError(t, err) + assert.True(t, result.Success) +} + +// Case 6: stage → delete → stage → delete cycle twice; final state has zero +// knowledge pack files and an empty sentinel. +func TestStageDeleteCycle(t *testing.T) { + ws := newTestWorkspace(t) + ctx := context.Background() + + for i := range 2 { + _, err := ws.StageKnowledgeFiles(ctx, &protocol.StageKnowledgeFilesArgs{ + Files: []protocol.KnowledgeFile{ + {RelPath: "cycle.md", Checksum: "cycle", ContentB64: b64("cycle")}, + }, + }) + require.NoError(t, err, "stage iteration %d", i) + + _, err = ws.DeleteKnowledgeFiles(ctx, &protocol.DeleteKnowledgeFilesArgs{ + RelPaths: []string{"cycle.md"}, + }) + require.NoError(t, err, "delete iteration %d", i) + } + + assert.NoFileExists(t, filepath.Join(ws.Root(), "cycle.md")) + + sentinel := readSentinelMap(t, ws.Root()) + assert.Empty(t, sentinel) +} + +// Case 7: path rejection rules. +func TestStageKnowledgeFiles_RejectedPaths(t *testing.T) { + ws := newTestWorkspace(t) + ctx := context.Background() + + cases := []struct { + relPath string + desc string + }{ + {"../etc/passwd", "path traversal with .."}, + {"sub/x.md", "subdirectory path with /"}, + {".hidden.md", "leading-dot filename"}, + {sentinelName, "sentinel filename itself"}, + {"", "empty string"}, + } + + for _, tc := range cases { + t.Run(tc.desc, func(t *testing.T) { + result, err := ws.StageKnowledgeFiles(ctx, &protocol.StageKnowledgeFilesArgs{ + Files: []protocol.KnowledgeFile{ + {RelPath: tc.relPath, Checksum: "sum", ContentB64: b64("data")}, + }, + }) + require.NoError(t, err) // handler never returns a hard error; rejection is in the per-file status + require.Len(t, result.Files, 1) + assert.False(t, result.Files[0].Success, "expected rejection for %q", tc.relPath) + assert.NotEmpty(t, result.Files[0].Error) + }) + } +} + +// TestValidateKnowledgeRelPath exercises the validation helper directly for +// thorough coverage of edge cases. +func TestValidateKnowledgeRelPath(t *testing.T) { + valid := []string{ + "DUTY.md", "runbook-api.md", "README.txt", "a", + } + for _, p := range valid { + assert.NoError(t, validateKnowledgeRelPath(p), "should be valid: %q", p) + } + + invalid := []string{ + "", "..", "a/b", `a\b`, ".hidden", sentinelName, "foo/../bar", + // These are now valid because ".." only matches the bare token, not + // substrings — "foo..bar" is a legitimate flat filename. + } + validButPreviouslyRejected := []string{"foo..bar", "v2..md"} + for _, p := range validButPreviouslyRejected { + assert.NoError(t, validateKnowledgeRelPath(p), "should be valid (double-dot in middle): %q", p) + } + for _, p := range invalid { + assert.Error(t, validateKnowledgeRelPath(p), "should be invalid: %q", p) + } +} diff --git a/ws/handler.go b/ws/handler.go index 95426e5..39c00bc 100644 --- a/ws/handler.go +++ b/ws/handler.go @@ -235,6 +235,22 @@ func (h *Handler) executeTask(ctx context.Context, req *protocol.TaskRequestPayl logger.Info("syncing skill", "skill", args.SkillName, "dir", args.SkillDir) return h.ws.SyncSkill(ctx, args) + case protocol.TaskOpStageKnowledgeFiles: + args, err := parseArgs[protocol.StageKnowledgeFilesArgs](req.Args) + if err != nil { + return nil, fmt.Errorf("invalid stage_knowledge_files args: %w", err) + } + logger.Info("staging knowledge files", "count", len(args.Files)) + return h.ws.StageKnowledgeFiles(ctx, args) + + case protocol.TaskOpDeleteKnowledgeFiles: + args, err := parseArgs[protocol.DeleteKnowledgeFilesArgs](req.Args) + if err != nil { + return nil, fmt.Errorf("invalid delete_knowledge_files args: %w", err) + } + logger.Info("deleting knowledge files", "count", len(args.RelPaths)) + return h.ws.DeleteKnowledgeFiles(ctx, args) + default: return nil, fmt.Errorf("unknown operation: %s", req.Operation) }