wehub-resource-sync
a06f331eb8
CI / benchmark (push) Has been skipped
install-script / posix-syntax (push) Successful in 6m1s
CI / build-onnx (push) Failing after 6m43s
init-smoke / dry-run (push) Failing after 15m57s
security / govulncheck (push) Has been cancelled
security / trivy-fs (push) Has been cancelled
CI / test (1.26, ubuntu-latest) (push) Has been cancelled
Scorecard supply-chain security / Scorecard analysis (push) Has been cancelled
CI / test (1.26, macos-latest) (push) Has been cancelled
CI / build-windows (push) Has been cancelled
CI / lint (push) Has been cancelled
install-script / powershell-syntax (push) Has been cancelled
install-script / install (macos-14) (push) Has been cancelled
install-script / install (ubuntu-latest) (push) Has been cancelled
296 lines
11 KiB
Go
296 lines
11 KiB
Go
package mcp
|
||
|
||
import (
|
||
"context"
|
||
"strings"
|
||
"testing"
|
||
|
||
"github.com/mark3labs/mcp-go/mcp"
|
||
"github.com/stretchr/testify/assert"
|
||
"github.com/stretchr/testify/require"
|
||
)
|
||
|
||
// TestDefaultFormatForClient verifies the GCX-capable client allowlist.
|
||
// The list is the contract: clients in it get gcx by default; everything
|
||
// else falls back to JSON. Drift in this list silently changes the wire
|
||
// format every shipping client receives, so the test pins it explicitly.
|
||
func TestDefaultFormatForClient(t *testing.T) {
|
||
cases := []struct {
|
||
client string
|
||
want string
|
||
}{
|
||
// GCX-capable: every client whose plugin/CLI ships a GCX1 decoder.
|
||
{"claude-code", "gcx"},
|
||
{"Claude-Code", "gcx"}, // case-insensitive
|
||
{" claude-code ", "gcx"}, // trimmed
|
||
{"cursor", "gcx"},
|
||
{"vscode", "gcx"},
|
||
{"zed", "gcx"},
|
||
{"aider", "gcx"},
|
||
{"kilocode", "gcx"},
|
||
{"opencode", "gcx"},
|
||
{"openclaw", "gcx"},
|
||
{"codex", "gcx"},
|
||
{"omp-coding-agent", "gcx"},
|
||
|
||
// Unknown / unset → JSON fallback.
|
||
{"", ""},
|
||
{"some-other-client", ""},
|
||
{"unknown", ""},
|
||
}
|
||
for _, tc := range cases {
|
||
assert.Equal(t, tc.want, defaultFormatForClient(tc.client),
|
||
"defaultFormatForClient(%q)", tc.client)
|
||
}
|
||
}
|
||
|
||
// TestResolveSessionFormat_NoSession returns "" for a bare context — no
|
||
// session means no client identity, which means no default format.
|
||
func TestResolveSessionFormat_NoSession(t *testing.T) {
|
||
srv, _ := setupTestServer(t)
|
||
assert.Equal(t, "", srv.resolveSessionFormat(context.Background()))
|
||
}
|
||
|
||
// TestResolveSessionFormat_KnownClient verifies the full per-session
|
||
// path: NoteSessionClient stores the client name, and resolveSessionFormat
|
||
// reads it back through sessionFor + defaultFormatForClient.
|
||
func TestResolveSessionFormat_KnownClient(t *testing.T) {
|
||
srv, _ := setupTestServer(t)
|
||
srv.NoteSessionClient("session_X", "claude-code", "1.0.42")
|
||
ctx := WithSessionID(context.Background(), "session_X")
|
||
assert.Equal(t, "gcx", srv.resolveSessionFormat(ctx))
|
||
}
|
||
|
||
func TestResolveSessionFormat_UnknownClient(t *testing.T) {
|
||
srv, _ := setupTestServer(t)
|
||
srv.NoteSessionClient("session_X", "some-bespoke-client", "0.1")
|
||
ctx := WithSessionID(context.Background(), "session_X")
|
||
assert.Equal(t, "", srv.resolveSessionFormat(ctx),
|
||
"unknown client must fall back to JSON (empty string)")
|
||
}
|
||
|
||
// TestNoteSessionClient_NilSafe ensures NoteSessionClient never panics
|
||
// when called on a nil *Server or with empty inputs — both are normal
|
||
// during boot races / embedded-mode tests.
|
||
func TestNoteSessionClient_NilSafe(t *testing.T) {
|
||
var srv *Server
|
||
srv.NoteSessionClient("sess", "claude-code", "1.0")
|
||
|
||
srv2, _ := setupTestServer(t)
|
||
srv2.NoteSessionClient("", "claude-code", "1.0") // empty session id → no-op
|
||
srv2.NoteSessionClient("sess", "", "1.0") // empty client → no-op
|
||
}
|
||
|
||
// TestNoteSessionClient_IsolatedPerSession verifies two sessions get
|
||
// independent client-name state. This is the core invariant that lets
|
||
// the daemon serve multiple proxies through one shared *Server.
|
||
func TestNoteSessionClient_IsolatedPerSession(t *testing.T) {
|
||
srv, _ := setupTestServer(t)
|
||
|
||
srv.NoteSessionClient("sess_A", "claude-code", "1.0")
|
||
srv.NoteSessionClient("sess_B", "some-bespoke-client", "0.1")
|
||
|
||
ctxA := WithSessionID(context.Background(), "sess_A")
|
||
ctxB := WithSessionID(context.Background(), "sess_B")
|
||
assert.Equal(t, "gcx", srv.resolveSessionFormat(ctxA))
|
||
assert.Equal(t, "", srv.resolveSessionFormat(ctxB))
|
||
}
|
||
|
||
// TestIsGCX_ExplicitFormatWins verifies that an explicit `format` arg
|
||
// overrides the per-session default in either direction.
|
||
func TestIsGCX_ExplicitFormatWins(t *testing.T) {
|
||
srv, _ := setupTestServer(t)
|
||
srv.NoteSessionClient("sess_A", "claude-code", "1.0") // session default = gcx
|
||
ctx := WithSessionID(context.Background(), "sess_A")
|
||
|
||
// Explicit "json" must override the session default.
|
||
req := mcp.CallToolRequest{}
|
||
req.Params.Arguments = map[string]any{"format": "json"}
|
||
assert.False(t, srv.isGCX(ctx, req),
|
||
"explicit format=json must defeat session-default gcx")
|
||
|
||
// Explicit "gcx" stays gcx.
|
||
req2 := mcp.CallToolRequest{}
|
||
req2.Params.Arguments = map[string]any{"format": "gcx"}
|
||
assert.True(t, srv.isGCX(ctx, req2))
|
||
}
|
||
|
||
// TestIsGCX_SessionDefaultApplies verifies that a session whose client
|
||
// is GCX-capable picks gcx when the request omits `format`.
|
||
func TestIsGCX_SessionDefaultApplies(t *testing.T) {
|
||
srv, _ := setupTestServer(t)
|
||
srv.NoteSessionClient("sess_A", "claude-code", "1.0")
|
||
ctx := WithSessionID(context.Background(), "sess_A")
|
||
|
||
req := mcp.CallToolRequest{}
|
||
req.Params.Arguments = map[string]any{} // no format
|
||
assert.True(t, srv.isGCX(ctx, req),
|
||
"claude-code session with no explicit format must default to gcx")
|
||
}
|
||
|
||
// TestIsGCX_NoSession_NoFormat returns false — the legacy default is
|
||
// JSON, and absent both an explicit format and a known client we must
|
||
// preserve that.
|
||
func TestIsGCX_NoSession_NoFormat(t *testing.T) {
|
||
srv, _ := setupTestServer(t)
|
||
req := mcp.CallToolRequest{}
|
||
req.Params.Arguments = map[string]any{}
|
||
assert.False(t, srv.isGCX(context.Background(), req))
|
||
}
|
||
|
||
// TestIsTOON_ExplicitFormatWins verifies that an explicit `format=toon`
|
||
// trips isTOON regardless of session default.
|
||
func TestIsTOON_ExplicitFormatWins(t *testing.T) {
|
||
srv, _ := setupTestServer(t)
|
||
srv.NoteSessionClient("sess_A", "claude-code", "1.0") // gcx by default
|
||
ctx := WithSessionID(context.Background(), "sess_A")
|
||
|
||
req := mcp.CallToolRequest{}
|
||
req.Params.Arguments = map[string]any{"format": "toon"}
|
||
assert.True(t, srv.isTOON(ctx, req))
|
||
assert.False(t, srv.isGCX(ctx, req),
|
||
"format=toon must not also trigger gcx")
|
||
}
|
||
|
||
// TestRespondJSONOrTOON_RoutesByFormat pins the helper that 14
|
||
// list-shaped tools share. With explicit format=toon the payload comes
|
||
// back as TOON-marshalled text; with no format and an unknown client
|
||
// it falls back to JSON. This is the single decision point most tools
|
||
// route through, so a regression here would silently flip every
|
||
// downstream consumer to the wrong format.
|
||
func TestRespondJSONOrTOON_RoutesByFormat(t *testing.T) {
|
||
srv, _ := setupTestServer(t)
|
||
payload := map[string]any{"x": 1, "y": "two"}
|
||
|
||
// format=toon → TOON text result.
|
||
reqTOON := mcp.CallToolRequest{}
|
||
reqTOON.Params.Arguments = map[string]any{"format": "toon"}
|
||
res, err := srv.respondJSONOrTOON(context.Background(), reqTOON, payload)
|
||
assert.NoError(t, err)
|
||
assert.False(t, res.IsError)
|
||
tc, ok := res.Content[0].(mcp.TextContent)
|
||
assert.True(t, ok, "expected TextContent for TOON result")
|
||
// TOON encodes scalar map values with `key: value` lines; an empty
|
||
// or JSON-shaped payload would not contain that exact prefix.
|
||
assert.Contains(t, tc.Text, "x: 1")
|
||
assert.NotContains(t, tc.Text, "{")
|
||
|
||
// no format, unknown client → JSON fallback.
|
||
reqJSON := mcp.CallToolRequest{}
|
||
reqJSON.Params.Arguments = map[string]any{}
|
||
res, err = srv.respondJSONOrTOON(context.Background(), reqJSON, payload)
|
||
assert.NoError(t, err)
|
||
tc, ok = res.Content[0].(mcp.TextContent)
|
||
assert.True(t, ok)
|
||
assert.Contains(t, tc.Text, "{") // JSON object braces
|
||
|
||
// format=json overrides session default → JSON.
|
||
srv.NoteSessionClient("sess_T", "claude-code", "1.0") // session default would be gcx, not toon
|
||
ctx := WithSessionID(context.Background(), "sess_T")
|
||
reqExplicitJSON := mcp.CallToolRequest{}
|
||
reqExplicitJSON.Params.Arguments = map[string]any{"format": "json"}
|
||
res, err = srv.respondJSONOrTOON(ctx, reqExplicitJSON, payload)
|
||
assert.NoError(t, err)
|
||
tc, ok = res.Content[0].(mcp.TextContent)
|
||
assert.True(t, ok)
|
||
assert.Contains(t, tc.Text, "{")
|
||
}
|
||
|
||
// TestRespondJSONOrTOON_MaxTokensTrimAndDecorate is the end-to-end
|
||
// pin for the new max_tokens parameter: a payload that exceeds the
|
||
// derived byte cap must come back trimmed AND carry the
|
||
// _truncated_by_tokens marker so the agent can see the tokens
|
||
// budget was the constraint that fired. Pairs with the unit tests
|
||
// in budget_test.go — this one rides the actual respondJSONOrTOON
|
||
// path that get_file_summary / find_usages / get_call_chain /
|
||
// get_editing_context / prefetch_context all share.
|
||
func TestRespondJSONOrTOON_MaxTokensTrimAndDecorate(t *testing.T) {
|
||
srv, _ := setupTestServer(t)
|
||
|
||
// Build a payload large enough that a tight max_tokens cap must
|
||
// trim. 200 rows × ~80 bytes per row ≈ 16 KB raw payload.
|
||
rows := make([]any, 200)
|
||
for i := range rows {
|
||
rows[i] = map[string]any{
|
||
"id": "row-" + strings.Repeat("x", 40),
|
||
"line": i,
|
||
"meta": strings.Repeat("padding-", 10),
|
||
}
|
||
}
|
||
payload := map[string]any{
|
||
"results": rows,
|
||
"total": 200,
|
||
}
|
||
|
||
// max_tokens=200 ≈ 700 bytes — clearly too small for the raw
|
||
// payload, so trim should fire.
|
||
req := mcp.CallToolRequest{}
|
||
req.Params.Arguments = map[string]any{
|
||
"format": "json",
|
||
"max_tokens": float64(200),
|
||
}
|
||
res, err := srv.respondJSONOrTOON(context.Background(), req, payload)
|
||
require.NoError(t, err)
|
||
require.False(t, res.IsError)
|
||
|
||
tc, ok := res.Content[0].(mcp.TextContent)
|
||
require.True(t, ok)
|
||
|
||
// The response must carry both the generic truncation flag AND
|
||
// the tokens-specific marker, plus the original max_tokens value
|
||
// for caller diagnostics.
|
||
assert.Contains(t, tc.Text, "_truncated_by_budget")
|
||
assert.Contains(t, tc.Text, "_truncated_by_tokens")
|
||
assert.Contains(t, tc.Text, "_max_tokens")
|
||
|
||
// Same payload with a generous max_tokens cap → no trim, no markers.
|
||
req.Params.Arguments = map[string]any{
|
||
"format": "json",
|
||
"max_tokens": float64(100000),
|
||
}
|
||
res, err = srv.respondJSONOrTOON(context.Background(), req, payload)
|
||
require.NoError(t, err)
|
||
tc, ok = res.Content[0].(mcp.TextContent)
|
||
require.True(t, ok)
|
||
assert.NotContains(t, tc.Text, "_truncated_by_budget")
|
||
assert.NotContains(t, tc.Text, "_truncated_by_tokens")
|
||
}
|
||
|
||
// TestGCXResponseWithBudget_MaxTokensTrimAndDecorate covers the GCX
|
||
// path's max_tokens behaviour: a GCX payload exceeding the
|
||
// token-derived byte cap gets row-trimmed by trimGCXBytes and then
|
||
// stamped with a `# max_tokens=N truncated_by_tokens=true` comment by
|
||
// decorateTokenBudgetGCX. The header stays intact so decoders still
|
||
// see a valid GCX1 envelope.
|
||
func TestGCXResponseWithBudget_MaxTokensTrimAndDecorate(t *testing.T) {
|
||
srv, _ := setupTestServer(t)
|
||
|
||
// Build a synthetic GCX payload that is comfortably over a
|
||
// 200-token (~700 byte) cap.
|
||
var sb strings.Builder
|
||
sb.WriteString("GCX1 tool=test fields=id,kind,name\n")
|
||
for i := 0; i < 50; i++ {
|
||
sb.WriteString("internal/foo.go::Bar")
|
||
sb.WriteString(strings.Repeat("x", 20))
|
||
sb.WriteString("\tfunction\tBar\n")
|
||
}
|
||
payload := []byte(sb.String())
|
||
|
||
req := mcp.CallToolRequest{}
|
||
req.Params.Arguments = map[string]any{
|
||
"format": "gcx",
|
||
"max_tokens": float64(200),
|
||
}
|
||
res, err := srv.gcxResponseWithBudget(req)(payload, nil)
|
||
require.NoError(t, err)
|
||
require.False(t, res.IsError)
|
||
|
||
tc, ok := res.Content[0].(mcp.TextContent)
|
||
require.True(t, ok)
|
||
assert.True(t, strings.HasPrefix(tc.Text, "GCX1 tool=test"))
|
||
// Both the byte-level and token-level markers must appear.
|
||
assert.Contains(t, tc.Text, "# truncated_by_budget=true")
|
||
assert.Contains(t, tc.Text, "max_tokens=200")
|
||
assert.Contains(t, tc.Text, "truncated_by_tokens=true")
|
||
}
|