From 47877c7e8bc2d4e850e7a7ece8889098a38e3b9e Mon Sep 17 00:00:00 2001 From: brook hong Date: Sun, 31 Aug 2025 19:11:04 +0800 Subject: [PATCH] fix: bedrock exception could be found at end of a stream (#2654) --- .gitignore | 1 - lua/avante/providers/bedrock.lua | 28 +++++++++++++--- lua/avante/utils/test.lua | 15 +++++++++ tests/data/bedrock_response_stream.bin | Bin 0 -> 6646 bytes ...bedrock_response_stream_with_exception.bin | Bin 0 -> 2275 bytes tests/providers/bedrock_spec.lua | 31 ++++++++++++++++++ 6 files changed, 69 insertions(+), 6 deletions(-) create mode 100644 lua/avante/utils/test.lua create mode 100644 tests/data/bedrock_response_stream.bin create mode 100644 tests/data/bedrock_response_stream_with_exception.bin diff --git a/.gitignore b/.gitignore index 039cd59..e65c79b 100644 --- a/.gitignore +++ b/.gitignore @@ -5,7 +5,6 @@ .venv __pycache__/ -data/ # Neovim plugin specific files plugin/packer_compiled.lua diff --git a/lua/avante/providers/bedrock.lua b/lua/avante/providers/bedrock.lua index eb4ef84..2faea64 100644 --- a/lua/avante/providers/bedrock.lua +++ b/lua/avante/providers/bedrock.lua @@ -79,6 +79,16 @@ function M:build_bedrock_payload(prompt_opts, request_body) return model_handler.build_bedrock_payload(self, prompt_opts, request_body) end +local function parse_exception(data) + local exceptions_found = {} + local bedrock_match = data:gmatch("exception(%b{})") + for bedrock_data_match in bedrock_match do + local jsn = vim.json.decode(bedrock_data_match) + table.insert(exceptions_found, "- " .. jsn.message) + end + return exceptions_found +end + function M:parse_stream_data(ctx, data, opts) -- @NOTE: Decode and process Bedrock response -- Each response contains a Base64-encoded `bytes` field, which is decoded into JSON. @@ -90,15 +100,23 @@ function M:parse_stream_data(ctx, data, opts) local json = vim.json.decode(data_stream) self:parse_response(ctx, data_stream, json.type, opts) end + local exceptions = parse_exception(data) + if #exceptions > 0 then + Utils.debug("Bedrock exceptions: ", vim.fn.json_encode(exceptions)) + if opts.on_chunk then + opts.on_chunk("\n**Exception caught**\n\n") + opts.on_chunk(table.concat(exceptions, "\n")) + end + vim.schedule(function() opts.on_stop({ reason = "error" }) end) + end end function M:parse_response_without_stream(data, event_state, opts) if opts.on_chunk == nil then return end - local bedrock_match = data:gmatch("exception(%b{})") - opts.on_chunk("\n**Exception caught**\n\n") - for bedrock_data_match in bedrock_match do - local jsn = vim.json.decode(bedrock_data_match) - opts.on_chunk("- " .. jsn.message .. "\n") + local exceptions = parse_exception(data) + if #exceptions > 0 then + opts.on_chunk("\n**Exception caught**\n\n") + opts.on_chunk(table.concat(exceptions, "\n")) end vim.schedule(function() opts.on_stop({ reason = "complete" }) end) end diff --git a/lua/avante/utils/test.lua b/lua/avante/utils/test.lua new file mode 100644 index 0000000..c273b69 --- /dev/null +++ b/lua/avante/utils/test.lua @@ -0,0 +1,15 @@ +-- This is a helper for unit tests. +local M = {} + +function M.read_file(fn) + fn = vim.uv.cwd() .. "/" .. fn + local file = io.open(fn, "r") + if file then + local data = file:read("*all") + file:close() + return data + end + return fn +end + +return M diff --git a/tests/data/bedrock_response_stream.bin b/tests/data/bedrock_response_stream.bin new file mode 100644 index 0000000000000000000000000000000000000000..91f9ad4bfbd9ac8c9b21766723f6b9325fc2e4f7 GIT binary patch literal 6646 zcmds+`)?Fg6vwv)f>DBqn(&jxA(CJ~W@oyEX+&k0eaz6Ew)>b@3~?U2JF`18yL4Zj zH5-Yd_=qB}1VVzMf`4AFMX({Uq{LmmT6YXt>X}T+*CeiIrBi{G;G4V`bxl^ZHKyoH z`b|fdr=3YhXqk*vR=r%HM-O+i#YVo6Y1B1VW#yzvTa3ksu3gk*al%rjor%PdXXIgR zU+5TVSl1D3ct*oybGG4!pKlpV&NUMJ81IQqRgr31p6*Q)On5HWHav-*DDd5EwFrO7 z81Hl)6|P%psSWQuCP!&lTciy(UN$|+Cg0(!jkYDvHsM_$Hm<9%qH8LWYj`5WzIZPh z#-HcMfJ?*h<54Tc(FW6U4L|ElR+E9)<9o11Rh7j$aH^8^t+d-QTyDtcj{-l@RvFt> zSix20Tv?Z6T}#Q8;T>i)Io1R2U0AOJ*H&JiEh|xu+D|oLn{>mpO67{}RBNv1*Bd8i znyq%HI~(+l#uLYqR3H6@KiP$@LayMSNWwVJdg$g;&X+~zDb6v{Tvd-Q&l|U`rKQma z@l_2bMnk-`!U#=+5odt;A>Pl$%!U1@HN~dMKEfW3l0D=dfh$iRIWjbUs*msDKhT*I zB;4yH+&$<{d&l~RThruOS|zFMJrwTXI!U9&@(M#saf5r1qB*ye3Cl^~{h?b;_4naF z@vrFfNMz&<5@%|96wPheICg7nni3VK{MGhm6~-<1#|%%Y0%MzmLS#NKe0OtG#U z3@s($_(WEmoSM!-+r5HrywCv}Xr3TJ7x9mirS*prN*dBWVN;$H-gh7`7yITckyVeA2?oP86&C7uq370!Yy~ORF*4GU94%ORNvY8J zOfE0VqMm1yOi>ovR;^W)WR=L#wkk(M$!mG8YXuAa+ssHbktI(qyGZX@Oon`+D9MVd z(NU_u6MfgiX8|z2O(<`{chKec*B|qe)L4=y1oPnfL-J&xOm$x*YIg7$kSTH7^m7eW zo-K#>oYmAB)qfFR-iNP1a&04F<9kyXe0np5?YB&V!%EvG$-LE63f0fzt9wH|a*%-7 zj=$eOvY9wIag%pA$z(ve1>Bg%R90p2=&A_unZOk85iI#~Q3K z^LW)U80uDXf^Ksi&2K~JHgSDrJ7hXA(udc~Cpc)5b)dkI;zAeExsRC5R9#-fw!CD= zWcf8IJZu51%$5heh+#)#V~41|rsCh<0c3j#vJresY-}oG*^VKnV0z>NO%Xd*{_jpl zVHXFvf9@4DM;iCYmN2#M;(6xf%>>F9s+t0@(wr}*S7~5Wzk@#AQvwLDg%Dt{h8jDw znJ_r9uKCYpf1c{gZ@rD@fY3{XkWB~Zv(1FZnOU3K$Jf(z2j-zQAAMYxo)v zA<8Fm+~d$qbYuOw3>k(}!-t`j(;al!GrODk0&MJd5&`-WoeJ=!ryRmE+kHuNdtiTO za`u8vy3)#iF!Q81u&aw?o>|q>F6Y5gRhH~~gOzA9=O}Ehqcc+54EV8Z;ppg4B}?^p zp|8;;{4MZ+b{jcf*wxCcAB@F)gkjAtfrOfn1y{=m9hGTJPbp{Zm1rPoh{2%^qKtKRY@bNWxMhfzuYZt&s z0<)-Eg5LxtkMTVbz7M7)m`=GwNy)j(?~B#!sF)dQFQT2r&qxIoOBluCKbj`nsxXwUff&2a)Wt%+IEvT|J`^8rzxW@&S?_AA7)A1j{Gdc6aUg_hwzjq zMeb*OoA(P<@GY64FrmaW)~?u%9Yb1sO1JPjC__=5xeYCF#1TD;Gm6atJ)&yvir(0c z^a9uSmwiKTIEprMqSSU9x@@nz3KnwYM9r4o9nfCq1u=$z9jgnX04`Mnh>?PCK;qs0 zoQOIj9n?aRmf4nZ>I{W-@5F_;(wszw&^tlPemH{fu!-oB+F?wP9$&H}kfORJ$V*D@H1G|2sGt-i~Dh+4o)N@BV0btQp_IU{kO9)*X6 zwnl5f2yIhKfTys+_Y}4uOnSlo1K>{!Y%yR~LuIz9kFnWYdX_$hGt3;bxdz^sfVTkK z^Y+!)s!RMmK7}%2;(Gpg643CJMRAf&XIVa9ELX+)K<+#`>?)6u`lQ#_j*bWV@ML80 zn^*Mm{C2JO(wmilM#smk|5!X%1?WnYC{l>)e-sg$JBqP)k!(oBrU1zcmZ)Y#)MikC z2PoE??r&lc94ip2qC!FUs)8b_O5?%9hu0*-v}}jp+}WOQpX~!1<&PC5+n>uX9`*iz zQcfLnCD@hvqd$>yX|I%bbt88KbfwERP_@qd$GaP-qWEvs!Eee>j^DitDtJ(({Nc3x zUcRvzqnIem3092E?Z{2i7)_jr?IFb|+(;_iLQV;l1+0uFL)b%v6 z_6OjtUh!6bTEA6(b7P59`$TjFOq)~Nass3{DKym<^_t_ilecOBgUe_xen<(D5-{pmaS;qojif4Tcc+h4hUT21qd0sdAV%#u83G@cw@ zF8#wdz^hl7Px*`20gRZFB;g_)7u?LB&iyRU?r|yeaprSNOmmL-<7DP@`K*A*9G+kr h-{