webfetch.test.ts 4.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154
  1. import { describe, expect, test } from "bun:test"
  2. import path from "path"
  3. import { Instance } from "../../src/project/instance"
  4. import { WebFetchTool } from "../../src/tool/webfetch"
  5. import { SessionID, MessageID } from "../../src/session/schema"
  6. const projectRoot = path.join(import.meta.dir, "../..")
  7. const ctx = {
  8. sessionID: SessionID.make("ses_test"),
  9. messageID: MessageID.make("message"),
  10. callID: "",
  11. agent: "build",
  12. abort: AbortSignal.any([]),
  13. messages: [],
  14. metadata: () => {},
  15. ask: async () => {},
  16. }
  17. type TimerID = ReturnType<typeof setTimeout>
  18. async function withFetch(
  19. mockFetch: (input: string | URL | Request, init?: RequestInit) => Promise<Response>,
  20. fn: () => Promise<void>,
  21. ) {
  22. const originalFetch = globalThis.fetch
  23. globalThis.fetch = mockFetch as unknown as typeof fetch
  24. try {
  25. await fn()
  26. } finally {
  27. globalThis.fetch = originalFetch
  28. }
  29. }
  30. async function withTimers(fn: (state: { ids: TimerID[]; cleared: TimerID[] }) => Promise<void>) {
  31. const set = globalThis.setTimeout
  32. const clear = globalThis.clearTimeout
  33. const ids: TimerID[] = []
  34. const cleared: TimerID[] = []
  35. globalThis.setTimeout = ((...args: Parameters<typeof setTimeout>) => {
  36. const id = set(...args)
  37. ids.push(id)
  38. return id
  39. }) as typeof setTimeout
  40. globalThis.clearTimeout = ((id?: TimerID) => {
  41. if (id !== undefined) cleared.push(id)
  42. return clear(id)
  43. }) as typeof clearTimeout
  44. try {
  45. await fn({ ids, cleared })
  46. } finally {
  47. ids.forEach(clear)
  48. globalThis.setTimeout = set
  49. globalThis.clearTimeout = clear
  50. }
  51. }
  52. describe("tool.webfetch", () => {
  53. test("returns image responses as file attachments", async () => {
  54. const bytes = new Uint8Array([137, 80, 78, 71, 13, 10, 26, 10])
  55. await withFetch(
  56. async () => new Response(bytes, { status: 200, headers: { "content-type": "IMAGE/PNG; charset=binary" } }),
  57. async () => {
  58. await Instance.provide({
  59. directory: projectRoot,
  60. fn: async () => {
  61. const webfetch = await WebFetchTool.init()
  62. const result = await webfetch.execute({ url: "https://example.com/image.png", format: "markdown" }, ctx)
  63. expect(result.output).toBe("Image fetched successfully")
  64. expect(result.attachments).toBeDefined()
  65. expect(result.attachments?.length).toBe(1)
  66. expect(result.attachments?.[0].type).toBe("file")
  67. expect(result.attachments?.[0].mime).toBe("image/png")
  68. expect(result.attachments?.[0].url.startsWith("data:image/png;base64,")).toBe(true)
  69. expect(result.attachments?.[0]).not.toHaveProperty("id")
  70. expect(result.attachments?.[0]).not.toHaveProperty("sessionID")
  71. expect(result.attachments?.[0]).not.toHaveProperty("messageID")
  72. },
  73. })
  74. },
  75. )
  76. })
  77. test("keeps svg as text output", async () => {
  78. const svg = '<svg xmlns="http://www.w3.org/2000/svg"><text>hello</text></svg>'
  79. await withFetch(
  80. async () =>
  81. new Response(svg, {
  82. status: 200,
  83. headers: { "content-type": "image/svg+xml; charset=UTF-8" },
  84. }),
  85. async () => {
  86. await Instance.provide({
  87. directory: projectRoot,
  88. fn: async () => {
  89. const webfetch = await WebFetchTool.init()
  90. const result = await webfetch.execute({ url: "https://example.com/image.svg", format: "html" }, ctx)
  91. expect(result.output).toContain("<svg")
  92. expect(result.attachments).toBeUndefined()
  93. },
  94. })
  95. },
  96. )
  97. })
  98. test("keeps text responses as text output", async () => {
  99. await withFetch(
  100. async () =>
  101. new Response("hello from webfetch", {
  102. status: 200,
  103. headers: { "content-type": "text/plain; charset=utf-8" },
  104. }),
  105. async () => {
  106. await Instance.provide({
  107. directory: projectRoot,
  108. fn: async () => {
  109. const webfetch = await WebFetchTool.init()
  110. const result = await webfetch.execute({ url: "https://example.com/file.txt", format: "text" }, ctx)
  111. expect(result.output).toBe("hello from webfetch")
  112. expect(result.attachments).toBeUndefined()
  113. },
  114. })
  115. },
  116. )
  117. })
  118. test("clears timeout when fetch rejects", async () => {
  119. await withTimers(async ({ ids, cleared }) => {
  120. await withFetch(
  121. async () => {
  122. throw new Error("boom")
  123. },
  124. async () => {
  125. await Instance.provide({
  126. directory: projectRoot,
  127. fn: async () => {
  128. const webfetch = await WebFetchTool.init()
  129. await expect(
  130. webfetch.execute({ url: "https://example.com/file.txt", format: "text" }, ctx),
  131. ).rejects.toThrow("boom")
  132. },
  133. })
  134. },
  135. )
  136. expect(ids).toHaveLength(1)
  137. expect(cleared).toHaveLength(1)
  138. expect(cleared[0]).toBe(ids[0])
  139. })
  140. })
  141. })