Skip to content

Commit 80e6099

Browse files
ivan-magdaivan-magda
committed
Add tool dispatch with read_file, write_file, edit_file
1 parent 053cd3a commit 80e6099

5 files changed

Lines changed: 559 additions & 34 deletions

File tree

‎Sources/Core/API/JSONValue.swift‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,16 @@ extension JSONValue {
6969
return nil
7070
}
7171

72+
public var intValue: Int? {
73+
if case .int(let value) = self {
74+
return value
75+
}
76+
if case .double(let value) = self {
77+
return Int(value)
78+
}
79+
return nil
80+
}
81+
7282
public subscript(key: String) -> JSONValue? {
7383
if case .object(let dict) = self {
7484
return dict[key]

‎Sources/Core/Agent.swift‎

Lines changed: 220 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ public final class Agent {
66
private let apiClient: any APIClientProtocol
77
private let model: String
88
private let systemPrompt: String
9+
private let workingDirectory: String
910
private let shellExecutor: ShellExecutor
1011
private var messages: [Message] = []
1112

@@ -18,6 +19,7 @@ public final class Agent {
1819
self.apiClient = apiClient
1920
self.model = model
2021
self.systemPrompt = systemPrompt ?? Self.buildSystemPrompt(cwd: workingDirectory)
22+
self.workingDirectory = workingDirectory
2123
self.shellExecutor = ShellExecutor(workingDirectory: workingDirectory)
2224
}
2325

@@ -32,7 +34,7 @@ public final class Agent {
3234
maxTokens: 4096,
3335
system: systemPrompt,
3436
messages: messages,
35-
tools: [Self.bashToolDefinition]
37+
tools: Self.toolDefinitions
3638
)
3739

3840
let response = try await apiClient.createMessage(request: request)
@@ -72,16 +74,11 @@ public final class Agent {
7274

7375
public static func buildSystemPrompt(cwd: String) -> String {
7476
"""
75-
You are a coding agent at \(cwd). You help the user by executing \
76-
shell commands to explore the filesystem, read and write files, run programs, \
77-
and accomplish tasks.
78-
79-
Guidelines:
80-
- Use the bash tool to execute commands
81-
- Always check the result of commands before proceeding
82-
- If a command fails, try to understand why and fix it
83-
- Be concise in your explanations
84-
- When editing files, show the relevant changes
77+
You are a coding agent at \(cwd). Use tools to solve tasks. \
78+
Act, don't explain.
79+
80+
- Prefer read_file/write_file/edit_file over bash for file operations
81+
- Always check tool results before proceeding
8582
"""
8683
}
8784
}
@@ -95,26 +92,99 @@ extension Agent {
9592
case executionFailed(String)
9693
}
9794

98-
private static let bashToolDefinition = ToolDefinition(
99-
name: "bash",
100-
description: "Run a shell command and return its output.",
101-
inputSchema: .object([
102-
"type": "object",
103-
"properties": .object([
104-
"command": .object([
105-
"type": "string",
106-
"description": "The shell command to execute"
107-
])
108-
]),
109-
"required": .array(["command"])
110-
])
111-
)
95+
static let toolDefinitions: [ToolDefinition] = [
96+
ToolDefinition(
97+
name: "bash",
98+
description: "Run a shell command.",
99+
inputSchema: .object([
100+
"type": "object",
101+
"properties": .object([
102+
"command": .object([
103+
"type": "string",
104+
"description": "The shell command to execute"
105+
])
106+
]),
107+
"required": .array(["command"])
108+
])
109+
),
110+
ToolDefinition(
111+
name: "read_file",
112+
description: "Read file contents.",
113+
inputSchema: .object([
114+
"type": "object",
115+
"properties": .object([
116+
"path": .object([
117+
"type": "string",
118+
"description": "The file path to read"
119+
]),
120+
"limit": .object([
121+
"type": "integer",
122+
"description": "Maximum number of lines to read"
123+
])
124+
]),
125+
"required": .array(["path"])
126+
])
127+
),
128+
ToolDefinition(
129+
name: "write_file",
130+
description: "Write content to a file.",
131+
inputSchema: .object([
132+
"type": "object",
133+
"properties": .object([
134+
"path": .object([
135+
"type": "string",
136+
"description": "The file path to write"
137+
]),
138+
"content": .object([
139+
"type": "string",
140+
"description": "The content to write"
141+
])
142+
]),
143+
"required": .array(["path", "content"])
144+
])
145+
),
146+
ToolDefinition(
147+
name: "edit_file",
148+
description: "Replace exact text in a file.",
149+
inputSchema: .object([
150+
"type": "object",
151+
"properties": .object([
152+
"path": .object([
153+
"type": "string",
154+
"description": "The file path to edit"
155+
]),
156+
"old_text": .object([
157+
"type": "string",
158+
"description": "The exact text to find and replace"
159+
]),
160+
"new_text": .object([
161+
"type": "string",
162+
"description": "The replacement text"
163+
])
164+
]),
165+
"required": .array(["path", "old_text", "new_text"])
166+
])
167+
)
168+
]
112169

113170
func executeTool(name: String, input: JSONValue) async -> Result<String, ToolError> {
114-
guard name == "bash" else {
171+
let handlers = [
172+
"bash": executeBash,
173+
"read_file": executeReadFile,
174+
"write_file": executeWriteFile,
175+
"edit_file": executeEditFile
176+
]
177+
178+
guard let handler = handlers[name] else {
115179
return .failure(.unknownTool(name))
116180
}
117181

182+
return await handler(input)
183+
}
184+
185+
// MARK: - Handlers
186+
187+
private func executeBash(_ input: JSONValue) async -> Result<String, ToolError> {
118188
guard let command = input["command"]?.stringValue else {
119189
return .failure(.missingParameter("command"))
120190
}
@@ -123,15 +193,136 @@ extension Agent {
123193
let result = try await shellExecutor.execute(command)
124194
return .success(result.formatted)
125195
} catch {
126-
return .failure(.executionFailed(error.localizedDescription))
196+
return .failure(.executionFailed("\(error)"))
127197
}
128198
}
129199

200+
private func executeReadFile(_ input: JSONValue) async -> Result<String, ToolError> {
201+
guard let path = input["path"]?.stringValue else {
202+
return .failure(.missingParameter("path"))
203+
}
204+
205+
switch resolveSafePath(path) {
206+
case .failure(let error):
207+
return .failure(error)
208+
case .success(let resolvedPath):
209+
do {
210+
let text = try String(contentsOfFile: resolvedPath, encoding: .utf8)
211+
212+
let lines = text.components(separatedBy: "\n")
213+
var output: String
214+
215+
if let limit = input["limit"]?.intValue, limit < lines.count {
216+
output =
217+
lines.prefix(limit).joined(separator: "\n")
218+
+ "\n... (\(lines.count - limit) more lines)"
219+
} else {
220+
output = text
221+
}
222+
223+
if output.count > 50_000 {
224+
output = String(output.prefix(50_000))
225+
}
226+
227+
return .success(output)
228+
} catch {
229+
return .failure(.executionFailed("\(error)"))
230+
}
231+
}
232+
}
233+
234+
private func executeWriteFile(_ input: JSONValue) async -> Result<String, ToolError> {
235+
guard let path = input["path"]?.stringValue else {
236+
return .failure(.missingParameter("path"))
237+
}
238+
239+
guard let content = input["content"]?.stringValue else {
240+
return .failure(.missingParameter("content"))
241+
}
242+
243+
switch resolveSafePath(path) {
244+
case .failure(let error):
245+
return .failure(error)
246+
case .success(let resolvedPath):
247+
do {
248+
let fileURL = URL(fileURLWithPath: resolvedPath)
249+
250+
try FileManager.default.createDirectory(
251+
at: fileURL.deletingLastPathComponent(),
252+
withIntermediateDirectories: true
253+
)
254+
try content.write(toFile: resolvedPath, atomically: true, encoding: .utf8)
255+
256+
return .success("Wrote \(content.utf8.count) bytes to \(path)")
257+
} catch {
258+
return .failure(.executionFailed("\(error)"))
259+
}
260+
}
261+
}
262+
263+
private func executeEditFile(_ input: JSONValue) async -> Result<String, ToolError> {
264+
guard let path = input["path"]?.stringValue else {
265+
return .failure(.missingParameter("path"))
266+
}
267+
268+
guard let oldText = input["old_text"]?.stringValue else {
269+
return .failure(.missingParameter("old_text"))
270+
}
271+
272+
guard let newText = input["new_text"]?.stringValue else {
273+
return .failure(.missingParameter("new_text"))
274+
}
275+
276+
switch resolveSafePath(path) {
277+
case .failure(let error):
278+
return .failure(error)
279+
case .success(let resolvedPath):
280+
do {
281+
var content = try String(contentsOfFile: resolvedPath, encoding: .utf8)
282+
283+
guard let range = content.range(of: oldText) else {
284+
return .failure(.executionFailed("Text not found in \(path)"))
285+
}
286+
287+
content.replaceSubrange(range, with: newText)
288+
try content.write(toFile: resolvedPath, atomically: true, encoding: .utf8)
289+
290+
return .success("Edited \(path)")
291+
} catch {
292+
return .failure(.executionFailed("\(error)"))
293+
}
294+
}
295+
}
296+
297+
// MARK: Helpers
298+
299+
private func resolveSafePath(_ relativePath: String) -> Result<String, ToolError> {
300+
let workDirURL = URL(fileURLWithPath: workingDirectory, isDirectory: true)
301+
let resolvedWorkDir = workDirURL.standardized
302+
303+
let fullURL =
304+
if relativePath.hasPrefix("/") {
305+
URL(fileURLWithPath: relativePath).standardized
306+
} else {
307+
workDirURL.appendingPathComponent(relativePath).standardized
308+
}
309+
310+
guard
311+
fullURL.path.hasPrefix(resolvedWorkDir.path + "/") || fullURL.path == resolvedWorkDir.path
312+
else {
313+
return .failure(.executionFailed("Path escapes workspace: \(relativePath)"))
314+
}
315+
316+
return .success(fullURL.path)
317+
}
318+
130319
private func printToolCall(name: String, input: JSONValue) {
131-
if let command = input["command"]?.stringValue {
320+
if name == "bash", let command = input["command"]?.stringValue {
132321
print("\(ANSIColor.yellow)$ \(command)\(ANSIColor.reset)")
322+
} else if let path = input["path"]?.stringValue {
323+
print("\(ANSIColor.yellow)> \(name): \(path)\(ANSIColor.reset)")
133324
} else {
134-
print("\(ANSIColor.yellow)⚡ \(name)\(ANSIColor.reset)")
325+
print("\(ANSIColor.yellow)> \(name)\(ANSIColor.reset)")
135326
}
136327
}
137328
}

0 commit comments

Comments
 (0)