diff --git a/Makefile b/Makefile index 7626834..e1baae4 100644 --- a/Makefile +++ b/Makefile @@ -12,6 +12,7 @@ test: $(LUA) tests/test_codemode.lua $(LUA) tests/test_provider_json.lua $(LUA) tests/test_client_lookup.lua + $(LUA) tests/test_call_tool_chain.lua $(LUA) tests/test_mcp.lua $(LUA) tests/test_template.lua $(LUA) tests/test_transports.lua diff --git a/lua-utcp-1.2-1.rockspec b/lua-utcp-1.2-1.rockspec index 9996886..596e820 100644 --- a/lua-utcp-1.2-1.rockspec +++ b/lua-utcp-1.2-1.rockspec @@ -8,6 +8,7 @@ build = { modules = { ["utcp"] = "lua/utcp/init.lua", ["utcp.client"] = "lua/utcp/client.lua", + ["utcp.chain"] = "lua/utcp/chain.lua", ["utcp.registry"] = "lua/utcp/registry.lua", ["utcp.errors"] = "lua/utcp/errors.lua", ["utcp.json"] = "lua/utcp/json.lua", diff --git a/lua/utcp/chain.lua b/lua/utcp/chain.lua new file mode 100644 index 0000000..61cfc7e --- /dev/null +++ b/lua/utcp/chain.lua @@ -0,0 +1,207 @@ +local Chain = {} + +local function resolve_arguments(arguments, state) + if type(arguments) == 'function' then + return arguments(state) + end + + if arguments == nil then + return {} + end + + return arguments +end + +local function normalize_error(step, err) + if type(err) == 'table' then + return err + end + + return { + step = step.name, + tool = step.tool, + error = tostring(err or 'tool call failed'), + } +end + +local function run_rollback(client, state, completed, errors) + local rollbacks = {} + + for i = #completed, 1, -1 do + local completed_step = completed[i] + local rollback = completed_step.rollback + + if rollback then + state.current = completed_step + + local rollback_tool = rollback.tool + local rollback_arguments = rollback.arguments + local ok, args_or_error = pcall(resolve_arguments, rollback_arguments, state) + + if not ok then + rollbacks[#rollbacks + 1] = { + step = completed_step.name, + tool = rollback_tool, + ok = false, + error = tostring(args_or_error), + } + else + local result, err = client:call_tool(rollback_tool, args_or_error or {}) + local rollback_result = { + step = completed_step.name, + tool = rollback_tool, + ok = err == nil, + output = result, + error = err, + } + + rollbacks[#rollbacks + 1] = rollback_result + + if err ~= nil then + errors[#errors + 1] = { + step = completed_step.name, + tool = rollback_tool, + error = err, + rollback = true, + } + end + end + end + end + + state.current = nil + return rollbacks +end + +function Chain.run(client, workflow) + if type(workflow) ~= 'table' then + return nil, 'call_tool_chain expects a workflow table' + end + + local state = { + ok = true, + steps = {}, + errors = {}, + previous = nil, + current = nil, + output = nil, + } + + local completed = {} + local failed_step + local failure_error + + for index, step in ipairs(workflow) do + if type(step) ~= 'table' then + return nil, { + failed_step = index, + error = 'workflow step must be a table', + steps = state.steps, + } + end + + local tool = step.tool or step.name + if type(tool) ~= 'string' or tool == '' then + return nil, { + failed_step = index, + error = 'workflow step requires a non-empty tool', + steps = state.steps, + } + end + + local name = step.name or tool + local record = { + index = index, + name = name, + tool = tool, + arguments = nil, + ok = false, + output = nil, + error = nil, + } + + state.current = record + + local args_ok, args_or_error = pcall(resolve_arguments, step.arguments, state) + if not args_ok then + record.error = tostring(args_or_error) + else + record.arguments = args_or_error or {} + + local call_ok, result_or_error, call_error = pcall( + client.call_tool, + client, + tool, + record.arguments + ) + + if not call_ok then + record.error = tostring(result_or_error) + elseif call_error ~= nil then + record.error = call_error + record.output = result_or_error + else + record.ok = true + record.output = result_or_error + end + end + + state.steps[#state.steps + 1] = record + state.steps[name] = record + state.previous = record + + if record.ok then + state.output = record.output + completed[#completed + 1] = { + index = index, + name = name, + tool = tool, + rollback = step.rollback, + output = record.output, + arguments = record.arguments, + } + else + state.ok = false + failed_step = index + failure_error = normalize_error(step, record.error) + state.errors[#state.errors + 1] = failure_error + + if workflow.on_error ~= 'continue' then + break + end + end + end + + state.current = nil + + if failed_step and workflow.on_error ~= 'continue' then + local result = { + ok = false, + failed_step = failed_step, + error = failure_error.error, + steps = state.steps, + errors = state.errors, + } + + if workflow.rollback then + result.rollbacks = run_rollback(client, state, completed, state.errors) + end + + return nil, result + end + + local result = { + ok = state.ok, + steps = state.steps, + errors = state.errors, + output = state.output, + } + + if failed_step then + result.failed_step = failed_step + end + + return result +end + +return Chain diff --git a/lua/utcp/client.lua b/lua/utcp/client.lua index a358484..ce907f7 100644 --- a/lua/utcp/client.lua +++ b/lua/utcp/client.lua @@ -1,4 +1,4 @@ -local json=require('utcp.json'); local Registry=require('utcp.registry'); local transports=require('utcp.transports'); local errors=require('utcp.errors') +local json=require('utcp.json'); local Registry=require('utcp.registry'); local transports=require('utcp.transports'); local errors=require('utcp.errors'); local Chain=require('utcp.chain') local Client={}; Client.__index=Client local aliases={streamable_http='streamable',streamable='streamable',http='http',sse='sse',tcp='tcp',udp='udp',cli='cli',text='text',graphql='graphql',mcp='mcp'} function Client.new(cfg) @@ -155,6 +155,9 @@ function Client:call_tool(name,args) args or {} ) end +function Client:call_tool_chain(workflow) + return Chain.run(self, workflow) +end function Client:call_tool_stream(name,args,on_event) local tool,p=self:find_tool(name); if not tool then return nil,p end; local tpl=tool.tool_call_template or tool.call_template or tool; local typ=aliases[tpl.call_template_type or tpl.provider_type or 'sse']; local cfg={}; if p then for k,v in pairs(p) do cfg[k]=v end end; for k,v in pairs(tpl) do cfg[k]=v end diff --git a/tests/test_call_tool_chain.lua b/tests/test_call_tool_chain.lua new file mode 100644 index 0000000..a157a56 --- /dev/null +++ b/tests/test_call_tool_chain.lua @@ -0,0 +1,122 @@ +package.path = './lua/?.lua;./lua/?/init.lua;' .. package.path + +local Client = require('utcp.client') + +local function assert_equal(actual, expected, message) + assert(actual == expected, message or ('expected ' .. tostring(expected) .. ', got ' .. tostring(actual))) +end + +local function new_client(calls) + local client = Client.new() + function client:call_tool(name, args) + calls[#calls + 1] = { name = name, args = args } + if name == 'read' then + return { text = 'hello' } + elseif name == 'transform' then + return { text = args.content .. ' world' } + elseif name == 'write' then + return { written = args.content } + elseif name == 'fail' then + return nil, 'boom' + elseif name == 'rollback' then + return { rolled_back = args.name } + end + return nil, 'unknown test tool: ' .. tostring(name) + end + return client +end + +-- Sequential execution passes the previous output into the next step. +do + local calls = {} + local client = new_client(calls) + local result, err = client:call_tool_chain({ + { name = 'read', tool = 'read', arguments = {} }, + { + name = 'transform', + tool = 'transform', + arguments = function(state) + return { content = state.previous.output.text } + end, + }, + { + name = 'write', + tool = 'write', + arguments = function(state) + return { content = state.steps.transform.output.text } + end, + }, + }) + + assert(not err, err) + assert(result.ok) + assert_equal(#result.steps, 3) + assert_equal(result.output.written, 'hello world') + assert_equal(calls[2].args.content, 'hello') + assert_equal(calls[3].args.content, 'hello world') +end + +-- A failed step stops the chain by default and records structured state. +do + local calls = {} + local client = new_client(calls) + local result, err = client:call_tool_chain({ + { name = 'read', tool = 'read', arguments = {} }, + { name = 'fail', tool = 'fail', arguments = {} }, + { name = 'write', tool = 'write', arguments = {} }, + }) + + assert(result == nil) + assert(err) + assert_equal(err.failed_step, 2) + assert_equal(err.error, 'boom') + assert_equal(#err.steps, 2) + assert_equal(#calls, 2) +end + +-- on_error=continue keeps executing while returning an unsuccessful result. +do + local calls = {} + local client = new_client(calls) + local result, err = client:call_tool_chain({ + on_error = 'continue', + { name = 'fail', tool = 'fail', arguments = {} }, + { name = 'write', tool = 'write', arguments = { content = 'after failure' } }, + }) + + assert(not err, err) + assert(not result.ok) + assert_equal(result.failed_step, 1) + assert_equal(#result.steps, 2) + assert_equal(result.output.written, 'after failure') +end + +-- rollback=true runs declared rollback tools in reverse order after failure. +do + local calls = {} + local client = new_client(calls) + local result, err = client:call_tool_chain({ + rollback = true, + { name = 'first', tool = 'read', arguments = {}, rollback = { + tool = 'rollback', + arguments = { name = 'first' }, + } }, + { name = 'second', tool = 'transform', arguments = { content = 'x' }, rollback = { + tool = 'rollback', + arguments = function(state) + return { name = state.current.name } + end, + } }, + { name = 'fail', tool = 'fail', arguments = {} }, + }) + + assert(result == nil) + assert(err) + assert_equal(#err.rollbacks, 2) + assert_equal(err.rollbacks[1].step, 'second') + assert_equal(err.rollbacks[2].step, 'first') + assert_equal(calls[4].args.name, 'second') + assert_equal(calls[5].args.name, 'first') +end + +print('call_tool_chain tests passed')