mistri 0.5.0 → 0.6.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (87) hide show
  1. checksums.yaml +4 -4
  2. data/CHANGELOG.md +469 -4
  3. data/CONTRIBUTING.md +52 -0
  4. data/README.md +289 -385
  5. data/SECURITY.md +40 -0
  6. data/UPGRADING.md +640 -0
  7. data/assets/logo-animated.svg +30 -0
  8. data/assets/logo-dark.svg +14 -0
  9. data/assets/logo-light.svg +14 -0
  10. data/assets/logo.svg +14 -0
  11. data/assets/social-preview.png +0 -0
  12. data/docs/README.md +87 -0
  13. data/docs/context-and-workspaces.md +378 -0
  14. data/docs/mcp.md +366 -0
  15. data/docs/reliability.md +450 -0
  16. data/docs/sessions.md +295 -0
  17. data/docs/sub-agents.md +401 -0
  18. data/docs/tool-contracts.md +324 -0
  19. data/examples/approval.rb +36 -0
  20. data/examples/browser.rb +27 -0
  21. data/examples/page_editor.rb +31 -0
  22. data/examples/quickstart.rb +21 -0
  23. data/lib/generators/mistri/install/install_generator.rb +7 -3
  24. data/lib/generators/mistri/install/templates/migration.rb.tt +2 -2
  25. data/lib/generators/mistri/mcp/templates/migration.rb.tt +1 -1
  26. data/lib/generators/mistri/mcp/templates/model.rb.tt +15 -8
  27. data/lib/mistri/agent.rb +575 -55
  28. data/lib/mistri/budget.rb +26 -1
  29. data/lib/mistri/child.rb +72 -16
  30. data/lib/mistri/compaction.rb +26 -10
  31. data/lib/mistri/compactor.rb +34 -11
  32. data/lib/mistri/console.rb +28 -7
  33. data/lib/mistri/content.rb +9 -3
  34. data/lib/mistri/dispatchers.rb +14 -12
  35. data/lib/mistri/errors.rb +83 -4
  36. data/lib/mistri/event.rb +24 -8
  37. data/lib/mistri/event_delivery.rb +60 -0
  38. data/lib/mistri/locks.rb +3 -3
  39. data/lib/mistri/mcp/client.rb +74 -19
  40. data/lib/mistri/mcp/egress.rb +216 -0
  41. data/lib/mistri/mcp/oauth.rb +476 -127
  42. data/lib/mistri/mcp/wires.rb +115 -23
  43. data/lib/mistri/mcp.rb +42 -8
  44. data/lib/mistri/message.rb +21 -11
  45. data/lib/mistri/models.rb +160 -22
  46. data/lib/mistri/providers/anthropic/assembler.rb +282 -44
  47. data/lib/mistri/providers/anthropic/serializer.rb +14 -9
  48. data/lib/mistri/providers/anthropic.rb +29 -6
  49. data/lib/mistri/providers/fake.rb +26 -10
  50. data/lib/mistri/providers/gemini/assembler.rb +148 -21
  51. data/lib/mistri/providers/gemini/serializer.rb +78 -9
  52. data/lib/mistri/providers/gemini.rb +31 -5
  53. data/lib/mistri/providers/openai/assembler.rb +337 -60
  54. data/lib/mistri/providers/openai/serializer.rb +13 -12
  55. data/lib/mistri/providers/openai.rb +29 -5
  56. data/lib/mistri/providers/schema_capabilities.rb +214 -0
  57. data/lib/mistri/result.rb +1 -1
  58. data/lib/mistri/schema.rb +893 -75
  59. data/lib/mistri/session.rb +560 -48
  60. data/lib/mistri/sinks/coalesced.rb +17 -10
  61. data/lib/mistri/spawner.rb +111 -61
  62. data/lib/mistri/sse.rb +57 -14
  63. data/lib/mistri/stores/active_record.rb +1 -1
  64. data/lib/mistri/stores/memory.rb +21 -2
  65. data/lib/mistri/sub_agent/execution.rb +81 -0
  66. data/lib/mistri/sub_agent/runtime.rb +297 -0
  67. data/lib/mistri/sub_agent.rb +124 -87
  68. data/lib/mistri/task_output.rb +24 -6
  69. data/lib/mistri/tool.rb +93 -13
  70. data/lib/mistri/tool_arguments.rb +377 -0
  71. data/lib/mistri/tool_call.rb +43 -9
  72. data/lib/mistri/tool_context.rb +4 -2
  73. data/lib/mistri/tool_executor.rb +117 -26
  74. data/lib/mistri/tool_result.rb +15 -10
  75. data/lib/mistri/tools/edit_file.rb +62 -8
  76. data/lib/mistri/tools.rb +41 -4
  77. data/lib/mistri/transport.rb +149 -44
  78. data/lib/mistri/usage.rb +65 -13
  79. data/lib/mistri/version.rb +1 -1
  80. data/lib/mistri/workspace/active_record.rb +183 -3
  81. data/lib/mistri/workspace/directory.rb +28 -8
  82. data/lib/mistri/workspace/memory.rb +34 -9
  83. data/lib/mistri/workspace/single.rb +62 -5
  84. data/lib/mistri/workspace.rb +39 -0
  85. data/lib/mistri.rb +6 -1
  86. data/mistri.gemspec +34 -0
  87. metadata +31 -3
@@ -8,77 +8,168 @@ module Mistri
8
8
  # concurrently up to max_concurrency; each runs inside the Rails executor
9
9
  # when Rails is present, so ActiveRecord connections return to the pool.
10
10
  #
11
- # A tool that raises becomes an in-band error string the model can read. An
12
- # abort never starts a not-yet-started call: it gets an interrupted result
13
- # instead, so the turn always pairs and the session replays cleanly.
11
+ # A tool that fails becomes an in-band ToolResult with an explicit error
12
+ # fact. An abort never starts a not-yet-started call: it gets an interrupted
13
+ # result instead, so the turn always pairs and the session replays cleanly.
14
14
  module ToolExecutor
15
+ # Separates Mistri's deadline from a handler's own Timeout::Error.
16
+ class InvocationTimeout < StandardError
17
+ end
18
+ private_constant :InvocationTimeout
19
+
15
20
  INTERRUPTED = "[interrupted: this tool call never ran]"
21
+ OUTCOME_UNKNOWN = "[interrupted: this tool call's outcome is unavailable; the tool may have " \
22
+ "executed, so verify its effects before retrying]"
23
+ VERIFY_BEFORE_RETRY = "The tool may have completed partially; verify its effects before " \
24
+ "retrying."
25
+ COMMITTED = Object.new.freeze
26
+ private_constant :COMMITTED
27
+
28
+ PreparedContext = Class.new(ToolContext) do
29
+ def arguments_prepared? = true
30
+ end
31
+ private_constant :PreparedContext
16
32
 
17
33
  module_function
18
34
 
19
35
  def call(calls, tools_by_name, signal: nil, max_concurrency: 4, session: nil, emit: nil,
20
36
  app: nil)
37
+ call_with_outcomes(
38
+ calls,
39
+ tools_by_name,
40
+ signal:,
41
+ max_concurrency:,
42
+ session:,
43
+ emit:,
44
+ app:,
45
+ prepared_arguments: false
46
+ ).map { |call, result, seconds, _committed| [call, result, seconds] }
47
+ end
48
+
49
+ # Agent prepares model arguments before approval and execution. This
50
+ # lower-level form preserves that boundary and exposes commitment so hooks
51
+ # cannot run for queued calls that an abort prevented from starting.
52
+ def call_with_outcomes(calls, tools_by_name, signal: nil, max_concurrency: 4, session: nil,
53
+ emit: nil, app: nil, prepared_arguments: false)
21
54
  return [] if calls.empty?
22
55
 
23
- context = ToolContext.new(session: session, signal: signal, emit: thread_safe(emit),
24
- app: app)
56
+ delivery = EventDelivery.wrap(emit, passthrough: [InvocationTimeout])
57
+ context_class = prepared_arguments ? PreparedContext : ToolContext
58
+ context = context_class.new(session: session, signal: signal,
59
+ emit: thread_safe(delivery, signal), app: app)
25
60
  results = Array.new(calls.length)
26
61
  queue = Queue.new
62
+ errors = Queue.new
27
63
  calls.each_with_index { |call, index| queue << [call, index] }
28
64
  workers = max_concurrency.clamp(1, calls.length)
29
- Array.new(workers) { worker(queue, results, tools_by_name, context) }.each(&:join)
30
- calls.zip(results).map do |call, entry|
31
- value, seconds = entry || [INTERRUPTED, nil]
32
- [call, value, seconds]
65
+ Array.new(workers) { worker(queue, results, tools_by_name, context, errors) }.each(&:join)
66
+ unless errors.empty?
67
+ error = errors.pop
68
+ raise EventDelivery.unwrap(error, delivery)
69
+ end
70
+
71
+ calls.each_with_index.map do |call, index|
72
+ entry = results[index]
73
+ entry = [failure(OUTCOME_UNKNOWN), nil, true] if entry.equal?(COMMITTED)
74
+ value, seconds, committed = entry || [failure(INTERRUPTED), nil, false]
75
+ [call, value, seconds, committed]
33
76
  end
34
77
  end
35
78
 
36
- def worker(queue, results, tools_by_name, context)
79
+ def worker(queue, results, tools_by_name, context, errors)
37
80
  Thread.new do
38
81
  loop do
82
+ break unless errors.empty?
83
+
39
84
  call, index = begin
40
85
  queue.pop(true)
41
86
  rescue ThreadError
42
87
  break
43
88
  end
44
89
  if context.signal&.aborted?
45
- results[index] = [INTERRUPTED, nil]
90
+ results[index] = [failure(INTERRUPTED), nil, false]
46
91
  next
47
92
  end
48
93
 
49
- started = Process.clock_gettime(Process::CLOCK_MONOTONIC)
50
- value = run_one(call, tools_by_name, context)
51
- results[index] = [value, Process.clock_gettime(Process::CLOCK_MONOTONIC) - started]
94
+ results[index] = run_one(call, index, results, tools_by_name, context)
52
95
  end
96
+ rescue StandardError => e
97
+ errors << e
53
98
  end
54
99
  end
55
100
 
56
- def run_one(call, tools_by_name, context)
101
+ def run_one(call, index, results, tools_by_name, context)
57
102
  tool = tools_by_name[call.name]
58
- return "Error: unknown tool #{call.name.inspect}" unless tool
103
+ return [failure("Error: unknown tool #{call.name.inspect}"), nil, false] unless tool
104
+
105
+ return [failure(INTERRUPTED), nil, false] unless commit(call, context)
106
+
107
+ results[index] = COMMITTED
108
+ [*invoke_one(tool, call, context), true]
109
+ end
59
110
 
60
- with_rails_executor { invoke(tool, call, context) }
111
+ def invoke_one(tool, call, context)
112
+ started = Process.clock_gettime(Process::CLOCK_MONOTONIC)
113
+ value = with_rails_executor { invoke(tool, call, context) }
114
+ [value, elapsed(started)]
115
+ rescue EventDelivery::Failure
116
+ raise
117
+ rescue InvocationTimeout
118
+ content = "Error running tool #{call.name.inspect}: timed out after #{tool.timeout}s. " \
119
+ "#{VERIFY_BEFORE_RETRY}"
120
+ [failure(content), elapsed(started)]
61
121
  rescue StandardError => e
62
- "Error running tool #{call.name.inspect}: #{e.class}: #{e.message}"
122
+ [failure("Error running tool #{call.name.inspect}: #{e.class}: #{e.message}. " \
123
+ "#{VERIFY_BEFORE_RETRY}"),
124
+ elapsed(started)]
63
125
  end
64
126
 
65
- # A tool with a timeout answers in band when it stalls, so one hung
66
- # handler cannot stall the whole run.
67
127
  def invoke(tool, call, context)
68
128
  return tool.call(call.arguments, context) unless tool.timeout
69
129
 
70
- Timeout.timeout(tool.timeout) { tool.call(call.arguments, context) }
71
- rescue Timeout::Error
72
- "Error running tool #{call.name.inspect}: timed out after #{tool.timeout}s"
130
+ Timeout.timeout(tool.timeout, InvocationTimeout) do
131
+ tool.call(call.arguments, context)
132
+ end
133
+ end
134
+
135
+ def commit(call, context)
136
+ return false if context.signal&.aborted?
137
+ return true unless context.emit
138
+
139
+ context.emit.call(Event.new(type: :tool_started, tool_call: call))
140
+ end
141
+
142
+ def failure(content) = ToolResult.new(content:, error: true)
143
+
144
+ def elapsed(started)
145
+ started && (Process.clock_gettime(Process::CLOCK_MONOTONIC) - started)
73
146
  end
74
147
 
75
148
  # Concurrent tools share the caller's sink; sinks are not required to be
76
149
  # thread-safe, so forwarded events serialize here.
77
- def thread_safe(emit)
78
- return nil unless emit
150
+ def thread_safe(delivery, signal)
151
+ return nil unless delivery
79
152
 
80
153
  mutex = Mutex.new
81
- ->(event) { mutex.synchronize { emit.call(event) } }
154
+ # Already-committed workers may report after a sibling fails delivery;
155
+ # give all of them the first failure without calling the broken sink.
156
+ failure = nil
157
+ lambda do |event|
158
+ mutex.synchronize do
159
+ next false if event.type == :tool_started && (failure || signal&.aborted?)
160
+ raise failure if failure
161
+
162
+ result = begin
163
+ delivery.call(event)
164
+ rescue InvocationTimeout
165
+ raise
166
+ rescue EventDelivery::Failure => e
167
+ failure = e
168
+ raise
169
+ end
170
+ event.type == :tool_started ? true : result
171
+ end
172
+ end
82
173
  end
83
174
 
84
175
  def with_rails_executor(&)
@@ -1,23 +1,28 @@
1
1
  # frozen_string_literal: true
2
2
 
3
3
  module Mistri
4
- # A two-channel tool result: content goes to the model, ui goes only to the
5
- # host. The ui payload rides the tool message and its :tool_result event,
6
- # persists with the session for transcript re-renders, and never reaches a
7
- # provider. Return one from a handler when the UI needs more than the model
8
- # should read or pay for: full query rows behind a compact answer, the
9
- # updated document behind "saved".
4
+ # A tool result with model content, host-only UI, and an explicit failure
5
+ # fact. Both channels and the error bit persist with the tool message; only
6
+ # content and provider-supported error signaling reach the model.
10
7
  #
11
- # Tool.define("edit_page", "Edits the page.") do |args|
12
- # page = apply(args)
8
+ # Tool.define("edit_page", "Edits the page.", schema: -> {
9
+ # object :changes, "Page changes", required: true
10
+ # }) do |args|
11
+ # page = apply(args.fetch("changes"))
13
12
  # Mistri::ToolResult.new(content: "Updated.", ui: { "html" => page })
14
13
  # end
15
14
  #
16
15
  # ui must be JSON-serializable; it is stored and delivered in canonical
17
16
  # JSON form (string keys), the same shape a reloaded session reads.
18
- ToolResult = Data.define(:content, :ui) do
19
- def initialize(content:, ui: nil)
17
+ ToolResult = Data.define(:content, :ui, :error) do
18
+ def initialize(content:, ui: nil, error: false)
19
+ unless [true, false].include?(error)
20
+ raise ArgumentError, "tool result error must be true or false"
21
+ end
22
+
20
23
  super
21
24
  end
25
+
26
+ def error? = error
22
27
  end
23
28
  end
@@ -2,6 +2,9 @@
2
2
 
3
3
  module Mistri
4
4
  module Tools
5
+ MAX_ATOMIC_EDIT_ATTEMPTS = 3
6
+ private_constant :MAX_ATOMIC_EDIT_ATTEMPTS
7
+
5
8
  module_function
6
9
 
7
10
  # The model-facing shape is flat {path, old_string, new_string,
@@ -10,28 +13,79 @@ module Mistri
10
13
  # back in band with the closest region and its exact difference, so the
11
14
  # model's retry is one shot.
12
15
  def edit_file(workspace)
16
+ atomic = atomic_workspace?(workspace)
13
17
  Tool.define("edit_file",
14
18
  "Replace an exact snippet of a document. Copy old_string verbatim from " \
15
19
  "read_file output including whitespace, without line-number prefixes. " \
16
20
  "It must match exactly one place; add surrounding lines to make it " \
17
21
  "unique, or set replace_all to change every occurrence.",
18
22
  eager_input_streaming: true,
23
+ argument_normalizer: Tools.method(:tolerate),
19
24
  schema: lambda {
20
25
  string :path, "Document path", required: true
21
26
  string :old_string, "Exact text to replace (whitespace matters)", required: true
22
27
  string :new_string, "Replacement text", required: true
23
28
  boolean :replace_all, "Replace every occurrence instead of exactly one"
24
29
  }) do |args|
25
- args = Tools.tolerate(args)
26
- with_document(workspace, args) do |content|
27
- result = Edit.replace(content, args["old_string"], args["new_string"],
28
- replace_all: args["replace_all"] == true)
29
- workspace.write(args["path"], result.content)
30
- "Replaced #{result.count} occurrence(s) in #{args["path"]}"
30
+ result = if atomic
31
+ replace_atomically(workspace, args)
32
+ else
33
+ replace_legacy(workspace, args)
34
+ end
35
+ next result if result.is_a?(ToolResult)
36
+
37
+ "Replaced #{result.count} occurrence(s) in #{args["path"]}"
38
+ rescue EditError, WorkspaceConflictError => e
39
+ ToolResult.new(content: "edit_file failed: #{e.message}", error: true)
40
+ end
41
+ end
42
+
43
+ def replace_legacy(workspace, args)
44
+ with_document(workspace, args) do |content|
45
+ result = replacement(content, args)
46
+ workspace.write(args["path"], result.content)
47
+ result
48
+ end
49
+ end
50
+
51
+ def replace_atomically(workspace, args)
52
+ MAX_ATOMIC_EDIT_ATTEMPTS.times do |attempt|
53
+ snapshot = workspace.snapshot(args["path"])
54
+ return missing_document(args["path"]) unless snapshot
55
+ unless snapshot.is_a?(Workspace::Snapshot)
56
+ raise TypeError, "workspace snapshot must be a Mistri::Workspace::Snapshot"
57
+ end
58
+
59
+ result = replacement(snapshot.content, args)
60
+ begin
61
+ committed = workspace.compare_and_write(
62
+ args["path"], result.content, expected_revision: snapshot.revision
63
+ )
64
+ unless committed.is_a?(Workspace::Snapshot)
65
+ raise TypeError,
66
+ "workspace compare_and_write must return a Mistri::Workspace::Snapshot"
67
+ end
68
+ unless same_content_bytes?(committed.content, result.content)
69
+ return ToolResult.new(
70
+ content: "The write to #{args["path"].inspect} committed, but storage " \
71
+ "transformed the resulting document. Use read_file before continuing.",
72
+ error: true
73
+ )
74
+ end
75
+ return result
76
+ rescue WorkspaceConflictError
77
+ raise if attempt == MAX_ATOMIC_EDIT_ATTEMPTS - 1
31
78
  end
32
- rescue EditError => e
33
- "edit_file failed: #{e.message}"
34
79
  end
35
80
  end
81
+
82
+ def replacement(content, args)
83
+ Edit.replace(content, args["old_string"], args["new_string"],
84
+ replace_all: args["replace_all"] == true)
85
+ end
86
+
87
+ def same_content_bytes?(left, right)
88
+ left == right || (left.bytesize == right.bytesize && left.b == right.b)
89
+ end
36
90
  end
37
91
  end
data/lib/mistri/tools.rb CHANGED
@@ -6,6 +6,8 @@ module Mistri
6
6
  # row in a documents table, or an actual file. The names stay read_file and
7
7
  # edit_file because those are the tool names models are trained on.
8
8
  module Tools
9
+ ATOMIC_WORKSPACE_METHODS = %i[snapshot compare_and_write].freeze
10
+ private_constant :ATOMIC_WORKSPACE_METHODS
9
11
  ALIASES = { "oldText" => "old_string", "old" => "old_string", "search" => "old_string",
10
12
  "newText" => "new_string", "new" => "new_string", "replace" => "new_string",
11
13
  "replaceAll" => "replace_all", "file" => "path", "filename" => "path" }.freeze
@@ -23,15 +25,50 @@ module Mistri
23
25
 
24
26
  def with_document(workspace, args)
25
27
  content = workspace.read(args["path"])
26
- return "No document at #{args["path"].inspect}. Use list_files to see paths." if content.nil?
28
+ return missing_document(args["path"]) if content.nil?
27
29
 
28
30
  yield content
29
31
  end
30
32
 
31
- # Absorb the drift real models produce: alias keys, stringly booleans,
32
- # unknown keys dropped by simply never being read.
33
+ def missing_document(path)
34
+ ToolResult.new(
35
+ content: "No document at #{path.inspect}. Use list_files to see paths.",
36
+ error: true
37
+ )
38
+ end
39
+
40
+ # Atomic writes are an explicit backend claim. A false or absent claim
41
+ # preserves the legacy four-method port; a true but incomplete claim fails
42
+ # before the model can rely on safety the backend does not implement.
43
+ def atomic_workspace?(workspace)
44
+ return false unless workspace.respond_to?(:atomic_writes?)
45
+
46
+ supported = workspace.atomic_writes?
47
+ unless [true, false].include?(supported)
48
+ raise ConfigurationError, "workspace atomic_writes? must return true or false"
49
+ end
50
+ return false unless supported
51
+
52
+ missing = ATOMIC_WORKSPACE_METHODS.reject { |method| workspace.respond_to?(method) }
53
+ unless missing.empty?
54
+ raise ConfigurationError,
55
+ "atomic workspace is missing #{missing.map(&:inspect).join(" and ")}"
56
+ end
57
+
58
+ true
59
+ end
60
+
61
+ # Absorb the drift real models produce for edit_file only. Ambiguous
62
+ # aliases fail instead of letting hash insertion order choose an edit.
33
63
  def tolerate(args)
34
- normalized = args.to_h { |key, value| [ALIASES.fetch(key.to_s, key.to_s), value] }
64
+ normalized = args.each_with_object({}) do |(key, value), copy|
65
+ canonical = ALIASES.fetch(key.to_s, key.to_s)
66
+ if copy.key?(canonical)
67
+ raise ArgumentError, "multiple arguments map to #{canonical.inspect}"
68
+ end
69
+
70
+ copy[canonical] = value
71
+ end
35
72
  case normalized["replace_all"]
36
73
  when "true", "1", 1 then normalized["replace_all"] = true
37
74
  when "false", "0", 0, nil then normalized["replace_all"] = false
@@ -1,6 +1,7 @@
1
1
  # frozen_string_literal: true
2
2
 
3
3
  require "net/http"
4
+ require "openssl"
4
5
  require "json"
5
6
  require "uri"
6
7
 
@@ -14,15 +15,66 @@ module Mistri
14
15
  # Streaming reads abort two ways: cooperatively between fragments, and hard,
15
16
  # by closing the socket from the abort signal's callback, so a stalled read
16
17
  # stops immediately instead of waiting out the read timeout.
18
+ # An address_resolver supplies a fresh validated set for each connection
19
+ # cycle; candidates are tried before the request is sent while the original
20
+ # hostname still owns Host, SNI, and TLS validation.
21
+ # JSON bodies and individual SSE lines share one configurable byte ceiling;
22
+ # a stream may contain any number of individually safe lines.
17
23
  class Transport
18
24
  KEEP_ALIVE_SECONDS = 30
25
+ ERROR_PREVIEW_BYTES = 500
26
+ BLANK_BODY = /\A[[:space:]]*\z/
27
+ CONNECT_ERRORS = [IOError, SocketError, SystemCallError, Timeout::Error,
28
+ Net::HTTPBadResponse, OpenSSL::SSL::SSLError].freeze
29
+
30
+ # Net::HTTP reconnects expired keep-alives internally. Resolving inside
31
+ # connect makes every MCP connection cycle cross the egress boundary.
32
+ class ResolvedHTTP < Net::HTTP
33
+ attr_writer :address_resolver
34
+
35
+ private
36
+
37
+ def connect
38
+ addresses = Array(@address_resolver.call)
39
+ raise ConfigurationError, "address_resolver returned no addresses" if addresses.empty?
40
+
41
+ original_timeout = @open_timeout
42
+ deadline = if original_timeout
43
+ Process.clock_gettime(Process::CLOCK_MONOTONIC) + original_timeout
44
+ end
45
+ failure = nil
46
+ addresses.each do |address|
47
+ remaining = deadline && (deadline - Process.clock_gettime(Process::CLOCK_MONOTONIC))
48
+ break if remaining && !remaining.positive?
49
+
50
+ @ipaddr = address
51
+ @open_timeout = remaining if remaining
52
+ begin
53
+ return super
54
+ rescue *CONNECT_ERRORS => e
55
+ failure = e
56
+ end
57
+ end
58
+ raise failure || Net::OpenTimeout.new("all approved addresses timed out")
59
+ ensure
60
+ @open_timeout = original_timeout if original_timeout
61
+ end
62
+ end
63
+ private_constant :ResolvedHTTP, :CONNECT_ERRORS, :BLANK_BODY
64
+
65
+ def initialize(origin:, open_timeout: 15, read_timeout: 300, write_timeout: 60,
66
+ address_resolver: nil, max_record_bytes: DEFAULT_MAX_RECORD_BYTES)
67
+ unless max_record_bytes.is_a?(Integer) && max_record_bytes.positive?
68
+ raise ConfigurationError, "max_record_bytes: must be a positive integer"
69
+ end
19
70
 
20
- def initialize(origin:, open_timeout: 15, read_timeout: 300, write_timeout: 60)
21
71
  @origin = origin.to_s.chomp("/")
22
72
  @uri = URI(@origin)
23
73
  @open_timeout = open_timeout
24
74
  @read_timeout = read_timeout
25
75
  @write_timeout = write_timeout
76
+ @address_resolver = address_resolver
77
+ @max_record_bytes = max_record_bytes
26
78
  @mutex = Mutex.new
27
79
  @connection = nil
28
80
  end
@@ -30,11 +82,19 @@ module Mistri
30
82
  # POST and decode a JSON response body. Retries once on a dead idle
31
83
  # socket, so it suits idempotent endpoints.
32
84
  def post(path, body:, headers: {})
33
- response = @mutex.synchronize do
34
- with_retry { connection.request(build_request(path, body, headers)) }
85
+ @mutex.synchronize do
86
+ with_retry do
87
+ parsed = nil
88
+ connection.request(build_request(path, body, headers)) do |response|
89
+ raise_for_status(response)
90
+ parsed = JSON.parse(read_json_body(response))
91
+ end
92
+ parsed
93
+ end
94
+ rescue ResponseTooLargeError, ResponseTooComplexError, ProviderError, JSON::ParserError
95
+ teardown
96
+ raise
35
97
  end
36
- raise_for_status(response)
37
- JSON.parse(response.body)
38
98
  end
39
99
 
40
100
  # POST and stream the SSE response, yielding each decoded data record.
@@ -47,32 +107,10 @@ module Mistri
47
107
 
48
108
  # POST for Streamable-HTTP endpoints (the MCP shape) that answer either
49
109
  # a JSON body or an SSE stream: yields each JSON record either way and
50
- # returns the response headers, downcased. Retries only a dead idle
51
- # socket that failed before any response started, so a side-effecting
52
- # call can never run twice.
53
- def post_either(path, body:, headers: {}, &block)
54
- @mutex.synchronize do
55
- retried = false
56
- begin
57
- started = false
58
- response_headers = nil
59
- connection.request(build_request(path, body, headers, streaming: true)) do |response|
60
- started = true
61
- raise_for_status(response)
62
- response_headers = response.to_hash.transform_values(&:first)
63
- read_either(response, &block)
64
- end
65
- response_headers
66
- rescue IOError, SocketError, SystemCallError, Timeout::Error => e
67
- teardown
68
- if started || retried || e.is_a?(Timeout::Error)
69
- raise ProviderError, "connection failed: #{e.message}"
70
- end
71
-
72
- retried = true
73
- retry
74
- end
75
- end
110
+ # returns the response headers, downcased. A replayable request retries
111
+ # once when a dead idle socket fails before any response starts.
112
+ def post_either(path, body:, headers: {}, replayable: true, &block)
113
+ @mutex.synchronize { post_either_locked(path, body, headers, replayable, &block) }
76
114
  end
77
115
 
78
116
  def close
@@ -81,14 +119,44 @@ module Mistri
81
119
 
82
120
  private
83
121
 
122
+ def post_either_locked(path, body, headers, replayable, &block)
123
+ retried = false
124
+ begin
125
+ started = false
126
+ response_headers = nil
127
+ connection.request(build_request(path, body, headers, streaming: true)) do |response|
128
+ started = true
129
+ raise_for_status(response)
130
+ response_headers = response.to_hash.transform_values(&:first)
131
+ read_either(response, &block)
132
+ end
133
+ response_headers
134
+ rescue ResponseTooLargeError, ResponseTooComplexError, ProviderError
135
+ teardown
136
+ raise
137
+ rescue IOError, SocketError, SystemCallError, Timeout::Error, JSON::ParserError,
138
+ Net::HTTPBadResponse, OpenSSL::SSL::SSLError => e
139
+ teardown
140
+ unless replayable
141
+ raise AmbiguousDeliveryError, "#{AmbiguousDeliveryError.default_message}: #{e.message}"
142
+ end
143
+ if started || retried || e.is_a?(Timeout::Error)
144
+ raise ProviderError, "connection failed: #{e.message}"
145
+ end
146
+
147
+ retried = true
148
+ retry
149
+ end
150
+ end
151
+
84
152
  def read_either(response, &block)
85
153
  if response["content-type"].to_s.include?("text/event-stream")
86
- sse = SSE.new
154
+ sse = SSE.new(max_record_bytes: @max_record_bytes)
87
155
  response.read_body { |chunk| sse.feed(chunk, &block) }
88
156
  sse.finish(&block)
89
157
  else
90
- raw = response.read_body
91
- block.call(JSON.parse(raw)) unless raw.to_s.strip.empty?
158
+ raw = read_json_body(response)
159
+ block.call(JSON.parse(raw)) unless BLANK_BODY.match?(raw)
92
160
  end
93
161
  end
94
162
 
@@ -115,6 +183,9 @@ module Mistri
115
183
  # than let the next request read stale frames.
116
184
  teardown if aborted
117
185
  aborted ? :aborted : nil
186
+ rescue EventDelivery::Failure, ResponseTooLargeError, ResponseTooComplexError, ProviderError
187
+ teardown
188
+ raise
118
189
  rescue IOError, SocketError, SystemCallError, Timeout::Error => e
119
190
  teardown
120
191
  return :aborted if signal&.aborted?
@@ -130,7 +201,7 @@ module Mistri
130
201
  end
131
202
 
132
203
  def read_stream(response, signal, &block)
133
- sse = SSE.new
204
+ sse = SSE.new(max_record_bytes: @max_record_bytes)
134
205
  aborted = false
135
206
  response.read_body do |fragment|
136
207
  if signal&.aborted?
@@ -146,17 +217,32 @@ module Mistri
146
217
  def build_request(path, body, headers, streaming: false)
147
218
  request = Net::HTTP::Post.new(URI("#{@origin}#{path}"))
148
219
  request["Content-Type"] = "application/json"
149
- if streaming
150
- request["Accept"] = "text/event-stream"
151
- # Net::HTTP silently negotiates gzip, and its inflater buffers the
152
- # whole stream, delivering "live" events in one burst at the end.
153
- request["Accept-Encoding"] = "identity"
154
- end
220
+ request["Accept"] = "text/event-stream" if streaming
155
221
  headers.each { |key, value| request[key] = value }
222
+ # Identity keeps the byte ceiling meaningful before an untrusted
223
+ # compressed expansion and preserves immediate SSE delivery.
224
+ request["Accept-Encoding"] = "identity"
156
225
  request.body = JSON.generate(body)
157
226
  request
158
227
  end
159
228
 
229
+ def read_json_body(response)
230
+ declared = response["content-length"]
231
+ if declared&.match?(/\A\d+\z/) && declared.to_i > @max_record_bytes
232
+ raise ResponseTooLargeError.new(kind: :json_body, limit: @max_record_bytes)
233
+ end
234
+
235
+ body = +""
236
+ response.read_body do |fragment|
237
+ if body.bytesize + fragment.bytesize > @max_record_bytes
238
+ raise ResponseTooLargeError.new(kind: :json_body, limit: @max_record_bytes)
239
+ end
240
+
241
+ body << fragment
242
+ end
243
+ body
244
+ end
245
+
160
246
  def with_retry
161
247
  attempted = false
162
248
  begin
@@ -176,13 +262,21 @@ module Mistri
176
262
  end
177
263
 
178
264
  def connection
179
- @connection ||= Net::HTTP.new(@uri.host, @uri.port).tap do |http|
265
+ @connection ||= begin
266
+ http = if @address_resolver
267
+ ResolvedHTTP.new(@uri.hostname, @uri.port, nil).tap do |resolved|
268
+ resolved.address_resolver = @address_resolver
269
+ end
270
+ else
271
+ Net::HTTP.new(@uri.hostname, @uri.port)
272
+ end
180
273
  http.use_ssl = @uri.scheme == "https"
181
274
  http.open_timeout = @open_timeout
182
275
  http.read_timeout = @read_timeout
183
276
  http.write_timeout = @write_timeout
184
277
  http.keep_alive_timeout = KEEP_ALIVE_SECONDS
185
278
  http.start
279
+ http
186
280
  end
187
281
  end
188
282
 
@@ -204,10 +298,21 @@ module Mistri
204
298
  status = response.code.to_i
205
299
  return if (200..299).cover?(status)
206
300
 
301
+ preview = +""
302
+ response.read_body do |fragment|
303
+ remaining = ERROR_PREVIEW_BYTES - preview.bytesize
304
+ preview << fragment.byteslice(0, remaining) if remaining.positive?
305
+ raise status_error(response, status, preview) if preview.bytesize >= ERROR_PREVIEW_BYTES
306
+ end
307
+ raise status_error(response, status, preview)
308
+ end
309
+
310
+ def status_error(response, status, preview)
207
311
  klass = error_class(status)
208
- options = { status: status, body: response.read_body.to_s[0, 500] }
312
+ body = preview.dup.force_encoding(Encoding::UTF_8).scrub("?")
313
+ options = { status: status, body: body }
209
314
  options[:retry_after] = retry_after(response) if klass == RateLimitError
210
- raise klass.new(**options)
315
+ klass.new(**options)
211
316
  end
212
317
 
213
318
  def error_class(status)