Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion go/adk/cmd/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -173,7 +173,7 @@ func main() {
logger.Info("Memory service enabled", "appName", appName)
}

runnerConfig, subagentSessionIDs, err := runnerpkg.CreateRunnerConfig(ctx, agentConfig, sessionService, appName, memoryService)
runnerConfig, subagentSessionIDs, err := runnerpkg.CreateRunnerConfig(ctx, agentConfig, sessionService, appName, memoryService, kagentURL, httpClient)
if err != nil {
logger.Error(err, "Failed to create Google ADK Runner config")
os.Exit(1)
Expand Down
21 changes: 21 additions & 0 deletions go/adk/pkg/runner/adapter.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ package runner
import (
"context"
"fmt"
"net/http"
"os"
"strings"

Expand All @@ -11,6 +12,7 @@ import (
kagentmemory "github.com/kagent-dev/kagent/go/adk/pkg/memory"
"github.com/kagent-dev/kagent/go/adk/pkg/session"
"github.com/kagent-dev/kagent/go/adk/pkg/sts"
"github.com/kagent-dev/kagent/go/adk/pkg/tools"
"github.com/kagent-dev/kagent/go/api/adk"
adkmemory "google.golang.org/adk/memory"
adkplugin "google.golang.org/adk/plugin"
Expand All @@ -34,6 +36,8 @@ func CreateRunnerConfig(
sessionService *session.KAgentSessionService,
appName string,
memoryService *kagentmemory.KagentMemoryService,
kagentURL string,
httpClient *http.Client,
) (runner.Config, map[string]string, error) {
log := logr.FromContextOrDiscard(ctx)

Expand All @@ -46,6 +50,23 @@ func CreateRunnerConfig(
extraTools = append(extraTools, saveTool)
}

if agentConfig.ShareTools != nil && *agentConfig.ShareTools && kagentURL != "" && httpClient != nil {
createTool, err := tools.NewCreateShareLinkTool(httpClient, kagentURL, appName)
if err != nil {
return runner.Config{}, nil, fmt.Errorf("failed to create create_share_link tool: %w", err)
}
listTool, err := tools.NewListShareLinksTool(httpClient, kagentURL, appName)
if err != nil {
return runner.Config{}, nil, fmt.Errorf("failed to create list_share_links tool: %w", err)
}
deleteTool, err := tools.NewDeleteShareLinkTool(httpClient, kagentURL, appName)
if err != nil {
return runner.Config{}, nil, fmt.Errorf("failed to create delete_share_link tool: %w", err)
}
extraTools = append(extraTools, createTool, listTool, deleteTool)
log.Info("Share link tools enabled")
}

stsPlugin, err := buildTokenPropagationPlugin(ctx, log)
if err != nil {
return runner.Config{}, nil, err
Expand Down
205 changes: 205 additions & 0 deletions go/adk/pkg/tools/share_tools.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,205 @@
package tools

import (
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"net/url"
"os"
"strings"

"google.golang.org/adk/tool"
"google.golang.org/adk/tool/functiontool"
)

// shareClient holds the dependencies for share link tools, captured at construction time.
type shareClient struct {
baseURL string
uiURL string // KAGENT_UI_URL, used to build full share URLs
appName string
httpClient *http.Client
}

// parseAppName converts a Python-identifier app_name back to (namespace, name).
// Format: "namespace__NS__agent_name" with hyphens encoded as underscores.
func parseAppName(appName string) (namespace, name string) {
parts := strings.SplitN(appName, "__NS__", 2)
if len(parts) != 2 {
return "", strings.ReplaceAll(appName, "_", "-")
}
return strings.ReplaceAll(parts[0], "_", "-"), strings.ReplaceAll(parts[1], "_", "-")
}

// shareURL returns the share URL for a session token.
// With uiURL set it returns a full absolute URL; otherwise a relative path.
func (c *shareClient) shareURL(token, sessionID string) string {
ns, name := parseAppName(c.appName)
path := fmt.Sprintf("/agents/%s/%s/chat/%s?share=%s", ns, name, sessionID, token)
if c.uiURL != "" {
return c.uiURL + path
}
return path
}

func (c *shareClient) do(ctx context.Context, method, path string) (*http.Response, error) {
req, err := http.NewRequestWithContext(ctx, method, c.baseURL+path, nil)
if err != nil {
return nil, fmt.Errorf("building request %s %s: %w", method, c.baseURL+path, err)
}
req.Header.Set("X-Agent-Name", c.appName)
return c.httpClient.Do(req)
}

func (c *shareClient) doWithJSON(ctx context.Context, method, path string, body io.Reader) (*http.Response, error) {
req, err := http.NewRequestWithContext(ctx, method, c.baseURL+path, body)
if err != nil {
return nil, fmt.Errorf("building request %s %s: %w", method, c.baseURL+path, err)
}
req.Header.Set("X-Agent-Name", c.appName)
req.Header.Set("Content-Type", "application/json")
return c.httpClient.Do(req)
}

func (c *shareClient) readBody(resp *http.Response) (map[string]any, error) {
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("reading response: %w", err)
}
var out map[string]any
if err := json.Unmarshal(body, &out); err != nil {
return nil, fmt.Errorf("decoding response: %w", err)
}
return out, nil
}

type createShareInput struct {
// ReadOnly controls whether the shared link allows visitors to send messages.
// When nil (not provided by the model), the server defaults to true (read-only).
ReadOnly *bool `json:"read_only,omitempty"`
}

// NewCreateShareLinkTool creates a tool that generates a share token for the current session.
func NewCreateShareLinkTool(httpClient *http.Client, baseURL, appName string) (tool.Tool, error) {
c := &shareClient{
baseURL: baseURL,
uiURL: strings.TrimRight(os.Getenv("KAGENT_UI_URL"), "/"),
appName: appName,
httpClient: httpClient,
}
return functiontool.New(functiontool.Config{
Name: "create_share_link",
Description: "Creates a share link for the current chat session. " +
"Returns a URL any authenticated user can open to view this session. " +
"The link is read-only by default (visitors cannot send messages). " +
"Set read_only=false to allow visitors to interact. " +
"Each call creates a new token; existing tokens remain valid.",
}, func(ctx tool.Context, in createShareInput) (map[string]any, error) {
sessionID := ctx.SessionID()
if sessionID == "" {
return nil, fmt.Errorf("create_share_link: no session ID in context")
}
reqBody, err := json.Marshal(in)
if err != nil {
return nil, fmt.Errorf("create_share_link: encoding request: %w", err)
}
resp, err := c.doWithJSON(ctx, http.MethodPost, "/api/sessions/"+url.PathEscape(sessionID)+"/shares", strings.NewReader(string(reqBody)))
if err != nil {
return nil, fmt.Errorf("create_share_link: request failed: %w", err)
}
if resp.StatusCode != http.StatusCreated {
_, _ = io.Copy(io.Discard, resp.Body)
_ = resp.Body.Close()
return nil, fmt.Errorf("create_share_link: unexpected status %d", resp.StatusCode)
}
body, err := c.readBody(resp)
if err != nil {
return nil, fmt.Errorf("create_share_link: %w", err)
}
data, _ := body["data"].(map[string]any)
token, _ := data["token"].(string)
readOnly, _ := data["read_only"].(bool)
return map[string]any{
"url": c.shareURL(token, sessionID),
"read_only": readOnly,
}, nil
})
}

// NewListShareLinksTool creates a tool that lists active share tokens for the current session.
func NewListShareLinksTool(httpClient *http.Client, baseURL, appName string) (tool.Tool, error) {
c := &shareClient{
baseURL: baseURL,
uiURL: strings.TrimRight(os.Getenv("KAGENT_UI_URL"), "/"),
appName: appName,
httpClient: httpClient,
}
return functiontool.New(functiontool.Config{
Name: "list_share_links",
Description: "Lists all active share links for the current session. " +
"Returns each share token and creation time. " +
"Use this to find a token before calling delete_share_link.",
}, func(ctx tool.Context, _ struct{}) (map[string]any, error) {
sessionID := ctx.SessionID()
if sessionID == "" {
return nil, fmt.Errorf("list_share_links: no session ID in context")
}
resp, err := c.do(ctx, http.MethodGet, "/api/sessions/"+url.PathEscape(sessionID)+"/shares")
if err != nil {
return nil, fmt.Errorf("list_share_links: request failed: %w", err)
}
if resp.StatusCode != http.StatusOK {
_, _ = io.Copy(io.Discard, resp.Body)
_ = resp.Body.Close()
return nil, fmt.Errorf("list_share_links: unexpected status %d", resp.StatusCode)
}
body, err := c.readBody(resp)
if err != nil {
return nil, fmt.Errorf("list_share_links: %w", err)
}
shares := body["data"]
if shares == nil {
shares = []any{}
}
return map[string]any{"shares": shares}, nil
})
}

type deleteShareInput struct {
Token string `json:"token"`
}

// NewDeleteShareLinkTool creates a tool that revokes a specific share token for the current session.
func NewDeleteShareLinkTool(httpClient *http.Client, baseURL, appName string) (tool.Tool, error) {
c := &shareClient{
baseURL: baseURL,
uiURL: strings.TrimRight(os.Getenv("KAGENT_UI_URL"), "/"),
appName: appName,
httpClient: httpClient,
}
return functiontool.New(functiontool.Config{
Name: "delete_share_link",
Description: "Deletes a share link by token, immediately revoking access for anyone using it. " +
"Use list_share_links first to find the token you want to revoke.",
}, func(ctx tool.Context, in deleteShareInput) (map[string]any, error) {
if in.Token == "" {
return nil, fmt.Errorf("delete_share_link: token is required")
}
sessionID := ctx.SessionID()
if sessionID == "" {
return nil, fmt.Errorf("delete_share_link: no session ID in context")
}
path := "/api/sessions/" + url.PathEscape(sessionID) + "/shares/" + url.PathEscape(in.Token)
resp, err := c.do(ctx, http.MethodDelete, path)
if err != nil {
return nil, fmt.Errorf("delete_share_link: request failed: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("delete_share_link: unexpected status %d", resp.StatusCode)
}
return map[string]any{"status": "revoked", "token": in.Token}, nil
})
}
Loading
Loading