diff --git a/agent/pkg/workflow/v3/pipeline.go b/agent/pkg/workflow/v3/pipeline.go index f795c6f45..e553950c0 100644 --- a/agent/pkg/workflow/v3/pipeline.go +++ b/agent/pkg/workflow/v3/pipeline.go @@ -3,16 +3,14 @@ package v3 import ( "context" "encoding/json" - "errors" "fmt" - "io" - "net/http" "regexp" "strings" "sync" "time" "github.com/malbeclabs/lake/agent/pkg/workflow" + "github.com/malbeclabs/lake/utils/pkg/docsfetch" ) // queryMarkerPattern matches [Q1], [Q2], etc. - markers that should only appear @@ -765,8 +763,9 @@ func formatCypherQueryResults(queries []CypherQueryInput, results []workflow.Exe return sb.String() } -// docsBaseURL is the base URL for fetching raw documentation from GitHub. -const docsBaseURL = "https://raw.githubusercontent.com/malbeclabs/docs/main/docs/" +// docsSource fetches documentation markdown. Shared with the MCP's read_docs so +// both paths apply one page bound and one slug validation. +var docsSource = docsfetch.FromEnv() // readDocs handles the read_docs tool - fetches documentation from GitHub. func (p *Workflow) readDocs(ctx context.Context, params map[string]any, onProgress workflow.ProgressCallback) (string, error) { @@ -785,23 +784,7 @@ func (p *Workflow) readDocs(ctx context.Context, params map[string]any, onProgre }) } - // Build URL and fetch - url := docsBaseURL + input.Page + ".md" - - req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) - if err != nil { - // Emit read_docs complete with error - if onProgress != nil { - onProgress(workflow.Progress{ - Stage: workflow.StageReadDocsComplete, - DocsPage: input.Page, - DocsError: err.Error(), - }) - } - return "", fmt.Errorf("failed to create request: %w", err) - } - - resp, err := http.DefaultClient.Do(req) + content, _, truncated, err := docsSource.Read(ctx, input.Page) if err != nil { // Emit read_docs complete with error if onProgress != nil { @@ -811,44 +794,10 @@ func (p *Workflow) readDocs(ctx context.Context, params map[string]any, onProgre DocsError: err.Error(), }) } - return "", fmt.Errorf("failed to fetch docs: %w", err) - } - defer resp.Body.Close() - - if resp.StatusCode != http.StatusOK { - errMsg := fmt.Sprintf("docs page not found: %s (status %d)", input.Page, resp.StatusCode) - // Emit read_docs complete with error - if onProgress != nil { - onProgress(workflow.Progress{ - Stage: workflow.StageReadDocsComplete, - DocsPage: input.Page, - DocsError: errMsg, - }) - } - return "", errors.New(errMsg) - } - - body, err := io.ReadAll(resp.Body) - if err != nil { - // Emit read_docs complete with error - if onProgress != nil { - onProgress(workflow.Progress{ - Stage: workflow.StageReadDocsComplete, - DocsPage: input.Page, - DocsError: err.Error(), - }) - } - return "", fmt.Errorf("failed to read response: %w", err) - } - - content := string(body) - - // Truncate if too long (docs shouldn't be huge, but just in case) - if len(content) > 10000 { - content = content[:10000] + "\n\n... (truncated)" + return "", err } - p.logInfo("workflow: docs fetched", "page", input.Page, "length", len(content)) + p.logInfo("workflow: docs fetched", "page", input.Page, "length", len(content), "truncated", truncated) // Emit read_docs complete if onProgress != nil { diff --git a/api/handlers/mcp.go b/api/handlers/mcp.go index 8dfc2d9fb..402ddae1c 100644 --- a/api/handlers/mcp.go +++ b/api/handlers/mcp.go @@ -5,6 +5,7 @@ import ( "encoding/json" "errors" "fmt" + "log/slog" "net/http" "reflect" "regexp" @@ -328,8 +329,10 @@ type ReadDocsInput struct { // ReadDocsOutput is the output from the read_docs tool. type ReadDocsOutput struct { - Page string `json:"page"` - Content string `json:"content"` + Page string `json:"page"` + Source string `json:"source,omitempty"` + Truncated bool `json:"truncated"` + Content string `json:"content"` } func (a *API) registerReadDocsTool(server *mcp.Server) { @@ -339,16 +342,22 @@ func (a *API) registerReadDocsTool(server *mcp.Server) { Description: "Read DoubleZero documentation to answer questions about concepts, architecture, setup, troubleshooting, or how the network works. Use this when users ask 'what is DZ', 'how do I set up', 'why isn't X working', or similar conceptual/procedural questions. Available pages include: index, architecture, setup, troubleshooting, connect, connect-multicast, contribute, contribute-overview, contribute-operations, users-overview, paying-fees, multicast-admin.", Annotations: &mcp.ToolAnnotations{ReadOnlyHint: true}, }, func(ctx context.Context, req *mcp.CallToolRequest, input ReadDocsInput) (*mcp.CallToolResult, ReadDocsOutput, error) { - // docsfetch validates the slug (path traversal) and truncates at its - // 10k page bound. - content, _, err := a.docsSource().Read(ctx, strings.TrimSpace(input.Page)) + // docsfetch validates the slug (path traversal) and bounds the page. + page := strings.TrimSpace(input.Page) + content, source, truncated, err := a.docsSource().Read(ctx, page) if err != nil { return nil, ReadDocsOutput{}, err } + if truncated { + slog.Warn("docs page truncated", "page", page, "source", source) + } + return nil, ReadDocsOutput{ - Page: strings.TrimSpace(input.Page), - Content: content, + Page: page, + Source: source, + Truncated: truncated, + Content: content, }, nil }) } diff --git a/api/handlers/mcp_onboarding.go b/api/handlers/mcp_onboarding.go index 3231a9940..b0b9809ba 100644 --- a/api/handlers/mcp_onboarding.go +++ b/api/handlers/mcp_onboarding.go @@ -5,6 +5,7 @@ import ( "encoding/json" "errors" "fmt" + "log/slog" "net" "net/http" "os/exec" @@ -74,10 +75,15 @@ var runbookListItem = regexp.MustCompile("(?m)^[\\t ]*[-*][\\t ]+(?:`([a-zA-Z0-9 var fencedBlock = regexp.MustCompile("(?s)```.*?```") func (a *API) loadRunbookCatalog(ctx context.Context) ([]runbookRef, error) { - content, _, err := a.docsSource().Read(ctx, runbookIndexPage) + content, _, truncated, err := a.docsSource().Read(ctx, runbookIndexPage) if err != nil { return nil, fmt.Errorf("failed to load runbook index: %w", err) } + // A short catalog would reject real runbooks as unknown, so refuse it outright. + if truncated { + slog.Warn("runbook index truncated; get_onboarding_runbook is down", "page", runbookIndexPage) + return nil, fmt.Errorf("runbook index %q exceeded the page bound; catalog is incomplete", runbookIndexPage) + } refs := parseRunbookIndex(content) if len(refs) == 0 { return nil, fmt.Errorf("runbook index %q has no runbook links", runbookIndexPage) @@ -173,6 +179,7 @@ type GetOnboardingRunbookOutput struct { Source string `json:"source,omitempty"` Mode string `json:"mode,omitempty"` AvailableRunbooks []string `json:"available_runbooks"` + Truncated bool `json:"truncated"` Preamble string `json:"preamble"` Runbook string `json:"runbook"` } @@ -203,10 +210,13 @@ func (a *API) registerGetOnboardingRunbookTool(server *mcp.Server) { return nil, GetOnboardingRunbookOutput{}, fmt.Errorf("unknown runbook %q (available: %s)", service, strings.Join(available, ", ")) } - content, source, err := a.loadRunbook(ctx, page) + content, source, truncated, err := a.docsSource().Read(ctx, page) if err != nil { return nil, GetOnboardingRunbookOutput{}, err } + if truncated { + slog.Warn("runbook truncated", "page", page, "source", source) + } return nil, GetOnboardingRunbookOutput{ Service: service, @@ -214,16 +224,13 @@ func (a *API) registerGetOnboardingRunbookTool(server *mcp.Server) { Source: source, Mode: strings.TrimSpace(input.Mode), AvailableRunbooks: available, + Truncated: truncated, Preamble: onboardingPreamble, Runbook: content, }, nil }) } -func (a *API) loadRunbook(ctx context.Context, page string) (content, source string, err error) { - return a.docsSource().Read(ctx, page) -} - // CheckEdgeAccessInput is the input for the check_edge_access tool. type CheckEdgeAccessInput struct { Pubkey string `json:"pubkey" jsonschema:"The user payer pubkey of the access pass (for self-service passes this is the wallet from 'doublezero address')"` diff --git a/api/handlers/mcp_onboarding_test.go b/api/handlers/mcp_onboarding_test.go index 605bf3b80..7f35b275a 100644 --- a/api/handlers/mcp_onboarding_test.go +++ b/api/handlers/mcp_onboarding_test.go @@ -85,6 +85,54 @@ func withRunbookDocs(t *testing.T, files map[string]string) *docsfetch.Client { } } +// A truncated runbook must say so on the output and in the content. +func TestMCPHandler_GetOnboardingRunbook_ReportsTruncation(t *testing.T) { + t.Parallel() + api := &handlers.API{DocsSource: withRunbookDocs(t, map[string]string{ + "runbooks.md": testRunbookIndex, + "feed-a-runbook.md": strings.Repeat("a", docsfetch.MaxPageBytes+1), + })} + handler, sessionID := mcpSession(t, api) + + output := callToolOutput(t, handler, sessionID, "get_onboarding_runbook", map[string]any{ + "service": "feed-a", + }) + assert.Equal(t, true, output["truncated"]) + assert.Contains(t, output["runbook"], "truncated at") +} + +func TestMCPHandler_ReadDocs_ReportsTruncation(t *testing.T) { + t.Parallel() + api := &handlers.API{DocsSource: withRunbookDocs(t, map[string]string{ + "big.md": strings.Repeat("a", docsfetch.MaxPageBytes+1), + "small.md": "# small", + })} + handler, sessionID := mcpSession(t, api) + + big := callToolOutput(t, handler, sessionID, "read_docs", map[string]any{"page": "big"}) + assert.Equal(t, true, big["truncated"]) + assert.Contains(t, big["content"], "truncated at") + assert.Contains(t, big["source"], "big.md") + + small := callToolOutput(t, handler, sessionID, "read_docs", map[string]any{"page": "small"}) + assert.Equal(t, false, small["truncated"]) + assert.Equal(t, "# small", small["content"]) +} + +func TestMCPHandler_GetOnboardingRunbook_RejectsTruncatedIndex(t *testing.T) { + t.Parallel() + api := &handlers.API{DocsSource: withRunbookDocs(t, map[string]string{ + "runbooks.md": testRunbookIndex + strings.Repeat("\n", docsfetch.MaxPageBytes), + })} + handler, sessionID := mcpSession(t, api) + + response := callTool(t, handler, sessionID, "get_onboarding_runbook", map[string]any{}) + result := response["result"].(map[string]any) + require.True(t, result["isError"].(bool), "expected isError, got: %v", result) + text := result["content"].([]any)[0].(map[string]any)["text"].(string) + assert.Contains(t, text, "catalog is incomplete") +} + func TestMCPHandler_GetOnboardingRunbook_ListsCatalog(t *testing.T) { t.Parallel() api := &handlers.API{DocsSource: withRunbookDocs(t, map[string]string{"runbooks.md": testRunbookIndex})} diff --git a/utils/pkg/docsfetch/docsfetch.go b/utils/pkg/docsfetch/docsfetch.go index 592fa55f0..6da6be841 100644 --- a/utils/pkg/docsfetch/docsfetch.go +++ b/utils/pkg/docsfetch/docsfetch.go @@ -16,12 +16,15 @@ import ( // DefaultBase is the public docs tree on malbeclabs/docs@main. const DefaultBase = "https://raw.githubusercontent.com/malbeclabs/docs/main/docs/" -// MaxPageBytes bounds a fetched page: content past this is dropped and marked -// truncated, so an oversized document cannot flood the model's context. -const MaxPageBytes = 10000 +// MaxPageBytes bounds a fetched page so an oversized document cannot flood the +// model's context; the largest docs page is ~34 KB. +const MaxPageBytes = 65536 -// truncationMarker is appended when a page exceeds MaxPageBytes. -const truncationMarker = "\n\n... (truncated)" +// truncationMarker tells the model the page was cut and where the rest is, so +// it can point the user at the source instead of answering from a partial page. +func truncationMarker(url string) string { + return fmt.Sprintf("\n\n... (truncated at %d bytes; full page: %s)", MaxPageBytes, url) +} // validPage restricts page slugs so callers cannot traverse out of the docs tree. var validPage = regexp.MustCompile(`^[a-zA-Z0-9][a-zA-Z0-9\-]*$`) @@ -53,14 +56,15 @@ func FromEnv() *Client { } } -// Read returns the markdown for page (without ".md") from GitHub raw. -func (c *Client) Read(ctx context.Context, page string) (content, source string, err error) { +// Read returns the markdown for page (without ".md") from GitHub raw. truncated +// reports whether the page exceeded MaxPageBytes and content is only its head. +func (c *Client) Read(ctx context.Context, page string) (content, source string, truncated bool, err error) { page = strings.TrimSpace(page) if page == "" { - return "", "", fmt.Errorf("page is required") + return "", "", false, fmt.Errorf("page is required") } if !ValidPage(page) { - return "", "", fmt.Errorf("invalid page name: %s", page) + return "", "", false, fmt.Errorf("invalid page name: %s", page) } base := c.Base @@ -75,22 +79,22 @@ func (c *Client) Read(ctx context.Context, page string) (content, source string, url := base + page + ".md" req, reqErr := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) if reqErr != nil { - return "", "", fmt.Errorf("failed to create request: %w", reqErr) + return "", "", false, fmt.Errorf("failed to create request: %w", reqErr) } resp, doErr := httpClient.Do(req) if doErr != nil { - return "", "", fmt.Errorf("failed to fetch docs: %w", doErr) + return "", "", false, fmt.Errorf("failed to fetch docs: %w", doErr) } defer resp.Body.Close() body, readErr := io.ReadAll(io.LimitReader(resp.Body, MaxPageBytes+1)) if readErr != nil { - return "", "", fmt.Errorf("failed to read response: %w", readErr) + return "", "", false, fmt.Errorf("failed to read response: %w", readErr) } if resp.StatusCode != http.StatusOK { - return "", "", fmt.Errorf("docs page not found: %s (status %d)", page, resp.StatusCode) + return "", "", false, fmt.Errorf("docs page not found: %s (status %d)", page, resp.StatusCode) } if len(body) > MaxPageBytes { - return string(body[:MaxPageBytes]) + truncationMarker, url, nil + return string(body[:MaxPageBytes]) + truncationMarker(url), url, true, nil } - return string(body), url, nil + return string(body), url, false, nil } diff --git a/utils/pkg/docsfetch/docsfetch_test.go b/utils/pkg/docsfetch/docsfetch_test.go index 51df86276..e15a3f49f 100644 --- a/utils/pkg/docsfetch/docsfetch_test.go +++ b/utils/pkg/docsfetch/docsfetch_test.go @@ -35,16 +35,17 @@ func TestRead_PublicDocs(t *testing.T) { HTTP: public.Client(), Base: public.URL + "/", } - content, source, err := c.Read(context.Background(), "edge-runbook") + content, source, truncated, err := c.Read(context.Background(), "edge-runbook") require.NoError(t, err) assert.Equal(t, "# from public", content) assert.Contains(t, source, "edge-runbook.md") + assert.False(t, truncated) } func TestRead_InvalidPage(t *testing.T) { t.Parallel() c := &Client{} - _, _, err := c.Read(context.Background(), "../../../etc/passwd") + _, _, _, err := c.Read(context.Background(), "../../../etc/passwd") require.Error(t, err) assert.Contains(t, err.Error(), "invalid page name") } @@ -58,10 +59,28 @@ func TestRead_TruncatesOversizedPage(t *testing.T) { t.Cleanup(public.Close) c := &Client{HTTP: public.Client(), Base: public.URL + "/"} - content, _, err := c.Read(context.Background(), "huge-page") + content, source, truncated, err := c.Read(context.Background(), "huge-page") require.NoError(t, err) - assert.Len(t, content, MaxPageBytes+len(truncationMarker)) - assert.True(t, strings.HasSuffix(content, truncationMarker)) + assert.True(t, truncated) + assert.Equal(t, big[:MaxPageBytes], strings.TrimSuffix(content, truncationMarker(source))) + assert.Contains(t, content, "truncated at 65536 bytes") + assert.Contains(t, content, "huge-page.md") +} + +func TestRead_LargestRealPageIsNotTruncated(t *testing.T) { + t.Parallel() + page := strings.Repeat("a", 34000) + public := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write([]byte(page)) + })) + t.Cleanup(public.Close) + + c := &Client{HTTP: public.Client(), Base: public.URL + "/"} + content, _, truncated, err := c.Read(context.Background(), "contribute-provisioning") + require.NoError(t, err) + assert.False(t, truncated) + assert.Equal(t, page, content) + assert.NotContains(t, content, "truncated") } func TestRead_NotFound(t *testing.T) { @@ -72,7 +91,7 @@ func TestRead_NotFound(t *testing.T) { t.Cleanup(public.Close) c := &Client{HTTP: public.Client(), Base: public.URL + "/"} - _, _, err := c.Read(context.Background(), "missing-page") + _, _, _, err := c.Read(context.Background(), "missing-page") require.Error(t, err) assert.Contains(t, err.Error(), "docs page not found") }