refactor(tools): use github.com/docker/aijson for tool-arg shape repair#2899
refactor(tools): use github.com/docker/aijson for tool-arg shape repair#2899trungutt wants to merge 1 commit into
Conversation
The four-repair shape-repair layer in pkg/tools/repair.go was extracted into github.com/docker/aijson (per the library's README). Swap the in-tree implementation for the upstream package so we get a single maintained copy and so other call sites in the codebase can adopt it without re-importing tools-internal symbols. Net change: - Add github.com/docker/aijson v0.1.0 - Replace the strict-then-repair dance in tools.NewHandler with a single aijson.Unmarshal call. Telemetry is preserved by wiring the existing tool_input_repaired slog event through aijson.OnRepair. - Delete pkg/tools/repair.go (269 lines) — equivalent logic now lives in the external package. - Rewrite pkg/tools/repair_test.go to exercise the same scenarios through the public NewHandler surface (which is what we actually rely on) rather than the now-deleted private helper. Behavior is unchanged: the same four repair kinds fire in the same order, and the same slog event is emitted with the same fields. Follow-up phases (separate PRs) will route the bypassing call sites (runtime/skill_runner.go, runtime/model_picker.go, ACP filesystem, background-agent handlers, etc.) through a shared helper so they also benefit from shape repair.
trungutt
force-pushed
the
feat/use-aijson-for-tool-args
branch
from
be39591 to
2fbcad8
Compare
| @@ -2,13 +2,18 @@ package tools | |||
|
|
|||
aheritier
added
area/tools
docker-agent
left a comment
docker-agent
left a comment
There was a problem hiding this comment.
Assessment: 🟢 APPROVE
This is a clean extraction refactor — the in-tree repair.go (~270 LOC) is deleted and replaced with a single call to aijson.Unmarshal from github.com/docker/aijson. The tools.go change is minimal and correct: the import swap is clean, the OnRepair callback preserves telemetry, and error propagation is maintained.
The test suite is sensibly migrated from white-box unit tests on tryRepairToolArgs to black-box integration tests through NewHandler/runHandler, covering all four repair kinds plus the multi-field, embedded-struct, and field-ordering cases.
One minor observation about test coverage is noted inline.
| require.Error(t, err) | ||
| } | ||
|
|
||
| func TestNewHandler_UnrepairableInputReturnsOriginalError(t *testing.T) { |
There was a problem hiding this comment.
[MINOR] Missing regression tests for non-object top-level input and unknown-field passthrough
Two defensive tests from the old repair_test.go were not ported as integration-level tests:
TestRepair_LeavesUnknownFieldsAlone— verifies that unknown fields in the JSON payload are not mutated by the repair pass.TestRepair_ReturnsFalseOnNonObjectInput— verifies that a top-level non-object payload (e.g."just a string"or42) does not panic or silently corrupt the result.
The non-object case is particularly relevant because LLMs occasionally send bare JSON strings as tool arguments. Without a regression test, if aijson.Unmarshal behaves differently on these inputs (panics, silently drops the payload, etc.), it won't be caught.
Consider adding integration-style tests via NewHandler/runHandler that cover both cases.
4 participants
What
Replace the in-tree tool-argument shape-repair layer (
pkg/tools/repair.go, ~270 LOC) with the upstreamgithub.com/docker/aijsonpackage, which is itself an extraction of that same code.Why
The four shape repairs (
unwrap_string_array,wrap_in_array,wrap_object_in_array,drop_null) were lifted into a standalone, reusable package per its README. Adopting the upstream copy means:tools.NewHandlerand roll their ownjson.Unmarshal(e.g.pkg/runtime/skill_runner.go,pkg/runtime/model_picker.go,pkg/acp/filesystem.go, background-agent handlers inpkg/tools/builtin/agent/agent.go). They'll be migrated in follow-up PRs without re-importing tools-internal symbols.edit_file's double-serializededitsbranch is subsumed byunwrap_string_array. A separate follow-up PR will simplifyParseEditFileArgsaccordingly (the byte-level syntax-repair layer for malformed JSON stays — aijson explicitly doesn't cover that).