Skip to content

Commit 68339e5

Browse files
committed
dex-analysis: add WARP diag analysis tools and reader D.O.
1 parent ce7ad70 commit 68339e5

File tree

6 files changed

+134
-0
lines changed

6 files changed

+134
-0
lines changed

apps/dex-analysis/README.md

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,8 @@ Currently available tools:
2020
| **Remote Captures** | `dex_list_remote_capture_eligible_devices` | Retrieve a list of devices eligible for remote captures like packet captures or WARP diagnostics. |
2121
| | `dex_create_remote_capture` | Initiate a remote capture on a specific device by id. |
2222
| | `dex_list_remote_captures` | Retrieve a list of previously created remote captures along with their details and status. |
23+
| | `dex_list_remote_warp_diag_contents` | List the filenames included in a remote WARP diag capture returned by `dex_list_remote_captures`. |
24+
| | `dex_explore_remote_warp_diag_output` | Retreive remote WARP diag file contents by filepath returned by `dex_list_remote_warp_diag_contents`. |
2325
| **Fleet Status** | `dex_fleet_status_live` | View live metrics for your fleet of zero trust devices for up to the past 1 hour. |
2426
| | `dex_fleet_status_over_time` | View historical metrics for your fleet of zero trust devices over time. |
2527
| | `dex_fleet_status_logs` | View historical logs for your fleet of zero trust devices for up to the past 7 days. |

apps/dex-analysis/src/dex-analysis.app.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ import type { AuthProps } from '@repo/mcp-common/src/cloudflare-oauth-handler'
1919
import type { Env } from './dex-analysis.context'
2020

2121
export { UserDetails }
22+
export { WarpDiagReader } from './warp_diag_reader'
2223

2324
const env = getEnv<Env>()
2425

apps/dex-analysis/src/dex-analysis.context.ts

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
import { WarpDiagReader } from './warp_diag_reader'
2+
13
import type { UserDetails } from '@repo/mcp-common/src/durable-objects/user_details.do'
24
import type { CloudflareDEXMCP } from './dex-analysis.app'
35

@@ -10,6 +12,7 @@ export interface Env {
1012
CLOUDFLARE_CLIENT_SECRET: string
1113
MCP_OBJECT: DurableObjectNamespace<CloudflareDEXMCP>
1214
USER_DETAILS: DurableObjectNamespace<UserDetails>
15+
WARP_DIAG_READER: DurableObjectNamespace<WarpDiagReader>
1316
MCP_METRICS: AnalyticsEngineDataset
1417
DEV_DISABLE_OAUTH: string
1518
DEV_CLOUDFLARE_API_TOKEN: string

apps/dex-analysis/src/tools/dex-analysis.tools.ts

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,18 @@
11
import { z } from 'zod'
22

33
import { fetchCloudflareApi } from '@repo/mcp-common/src/cloudflare-api'
4+
import { getEnv } from '@repo/mcp-common/src/env'
5+
6+
import { Env } from '../dex-analysis.context'
7+
import { getReader } from '../warp_diag_reader'
48

59
import type { ToolCallback } from '@modelcontextprotocol/sdk/server/mcp.js'
610
import type { ToolAnnotations } from '@modelcontextprotocol/sdk/types.js'
711
import type { ZodRawShape, ZodTypeAny } from 'zod'
812
import type { CloudflareDEXMCP } from '../dex-analysis.app'
913

14+
const env = getEnv<Env>()
15+
1016
export function registerDEXTools(agent: CloudflareDEXMCP) {
1117
registerTool({
1218
name: 'dex_test_statistics',
@@ -478,6 +484,47 @@ export function registerDEXTools(agent: CloudflareDEXMCP) {
478484
})
479485
},
480486
})
487+
488+
registerTool({
489+
name: 'dex_list_remote_warp_diag_contents',
490+
description:
491+
'Given a WARP diag remote capture download url, returns a list of the files contained in the archive.',
492+
schema: {
493+
download: z
494+
.string()
495+
.describe(
496+
'The `filename` url from the dex_list_remote_captures response for successful WARP diag captures.'
497+
),
498+
},
499+
llmContext:
500+
'Use the dex_explore_remote_warp_diag_output tool for specific file paths to explore the file contents for analysis.',
501+
agent,
502+
callback: async ({ accessToken, download }) => {
503+
const reader = await getReader(env, accessToken, download)
504+
return await reader.list(accessToken, download)
505+
},
506+
})
507+
508+
registerTool({
509+
name: 'dex_explore_remote_warp_diag_output',
510+
description:
511+
'Explore the contents of remote capture WARP diag archive filepaths returned by the dex_list_remote_warp_diag_contents tool for analysis.',
512+
schema: {
513+
download: z
514+
.string()
515+
.describe(
516+
'The `filename` url from the dex_list_remote_captures response for successful WARP diag captures.'
517+
),
518+
filepath: z.string().describe('The file path from the archive to retrieve contents for.'),
519+
},
520+
llmContext:
521+
'To avoid hitting conversation and memory limits, avoid outputting the whole contents of these files to the user unless specifically asked to. Instead prefer to show relevant snippets only.',
522+
agent,
523+
callback: async ({ accessToken, download, filepath }) => {
524+
const reader = await getReader(env, accessToken, download)
525+
return await reader.read(accessToken, download, filepath)
526+
},
527+
})
481528
}
482529

483530
// Helper to simplify tool registration by reducing boilerplate for accountId and accessToken
Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
import { DurableObject } from 'cloudflare:workers'
2+
import JSZip from 'jszip'
3+
4+
import { Env } from './dex-analysis.context'
5+
6+
// Helper for reading large WARP diag zip archives.
7+
// Holds the contents in memory between requests from the agent for specific files
8+
// instead of having the worker download the zip on every request.
9+
//
10+
// Each DO represents one remote capture zip
11+
export class WarpDiagReader extends DurableObject<Env> {
12+
#cache?: { files: string[]; zip: JSZip }
13+
14+
// List the files in the zip for the agent
15+
async list(accessToken: string, url: string) {
16+
const { files } = await this.#getZip(accessToken, url)
17+
return files
18+
}
19+
20+
// Return the contents of a file by path
21+
async read(accessToken: string, url: string, filepath: string) {
22+
const { zip } = await this.#getZip(accessToken, url)
23+
const f = zip.file(filepath)
24+
const content = await f?.async('text')
25+
return content
26+
}
27+
28+
async #getZip(accessToken: string, url: string) {
29+
if (this.#cache) {
30+
return this.#cache
31+
}
32+
33+
const res = await fetch(url, {
34+
headers: {
35+
Authorization: `Bearer ${accessToken}`,
36+
},
37+
})
38+
const zip = await new JSZip().loadAsync(await res.arrayBuffer())
39+
const files: string[] = []
40+
for (const [relativePath, file] of Object.entries(zip.files)) {
41+
if (!file.dir) {
42+
files.push(relativePath)
43+
}
44+
}
45+
46+
const cache = { files, zip }
47+
this.#cache = cache
48+
return cache
49+
}
50+
}
51+
52+
async function hashToken(accessToken: string) {
53+
const hashArr = Array.from(
54+
new Uint8Array(await crypto.subtle.digest('SHA-256', new TextEncoder().encode(accessToken)))
55+
)
56+
return hashArr.map((b) => b.toString(16).padStart(2, '0')).join('')
57+
}
58+
59+
// Create unique name based on accessToken hash and download url. In order to read cached zip from memory
60+
// you need to have the same access token that was used to fetch it.
61+
async function readerName(accessToken: string, url: string) {
62+
return (await hashToken(accessToken)) + url
63+
}
64+
65+
export async function getReader(env: Env, accessToken: string, download: string) {
66+
const name = await readerName(accessToken, download)
67+
const id = env.WARP_DIAG_READER.idFromName(name)
68+
return env.WARP_DIAG_READER.get(id)
69+
}

apps/dex-analysis/wrangler.jsonc

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,10 @@
2525
{
2626
"class_name": "UserDetails",
2727
"name": "USER_DETAILS"
28+
},
29+
{
30+
"class_name": "WarpDiagReader",
31+
"name": "WARP_DIAG_READER"
2832
}
2933
]
3034
},
@@ -65,6 +69,10 @@
6569
"class_name": "UserDetails",
6670
"name": "USER_DETAILS",
6771
"script_name": "mcp-cloudflare-workers-observability-staging"
72+
},
73+
{
74+
"class_name": "WarpDiagReader",
75+
"name": "WARP_DIAG_READER"
6876
}
6977
]
7078
},
@@ -98,6 +106,10 @@
98106
"class_name": "UserDetails",
99107
"name": "USER_DETAILS",
100108
"script_name": "mcp-cloudflare-workers-observability-production"
109+
},
110+
{
111+
"class_name": "WarpDiagReader",
112+
"name": "WARP_DIAG_READER"
101113
}
102114
]
103115
},

0 commit comments

Comments
 (0)